mstar.model.components.qwen3_lm

mstar.model.components.qwen3_lm#

Dense Qwen3 causal language model from the shared components.

Qwen3 = Llama-style pre-norm decoder with GQA, per-head QK RMSNorm, SwiGLU MLP and plain RoPE (rope_theta 1e6, no Llama-3 scaling). Several audio-LLM checkpoints ship one as their text decoder (Qwen3-ASR, Higgs-Audio, Qwen3-Omni’s dense layers); this is the one place to build it from ParallelAttention / ParallelGatedMLP / RMSNorm, bound to the KV, attention and position resources of the node that owns it.

Parameter paths mirror HF’s Qwen3ForCausalLM with the model. prefix dropped (embed_tokens.*, layers.N.*, norm.*, lm_head.*); a checkpoint that ties lm_head to the embedding loads without the head key.

Classes

Qwen3DenseLM(config, *, attn_key, kv_key, ...)

Qwen3LMConfig([hidden_size, ...])

class mstar.model.components.qwen3_lm.Qwen3DenseLM(config, *, attn_key, kv_key, pos_key, comm_group=None)[source]#

Bases: Module

Parameters:
embed(input_ids)[source]#
Parameters:

input_ids (Tensor)

Return type:

Tensor

forward(input_embeds, *, label)[source]#

(tokens, hidden) -> (tokens, hidden) after the final norm.

The label and layer index are cursors on the shared resources: bind the label once, advance the index per layer. Passing them as arguments instead would make inductor specialize on the int.

Parameters:
Return type:

Tensor

load_weights(weights)[source]#

Stream HF (name, tensor) pairs (model. prefix already dropped); tie the head to the embedding when the checkpoint does; raise on anything left unloaded.

Return type:

set[str]

logits(hidden_states)[source]#
Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.qwen3_lm.Qwen3LMConfig(hidden_size: 'int' = 2048, num_hidden_layers: 'int' = 28, num_attention_heads: 'int' = 16, num_key_value_heads: 'int' = 8, head_dim: 'int' = 128, intermediate_size: 'int' = 6144, vocab_size: 'int' = 151936, rms_norm_eps: 'float' = 1e-06, rope_theta: 'float' = 1000000.0, max_position_embeddings: 'int' = 32768, tie_word_embeddings: 'bool' = False)[source]#

Bases: object

Parameters:

  • hidden_size (int)

  • num_hidden_layers (int)

  • num_attention_heads (int)

  • num_key_value_heads (int)

  • head_dim (int)

  • intermediate_size (int)

  • vocab_size (int)

  • rms_norm_eps (float)

  • rope_theta (float)

  • max_position_embeddings (int)

  • tie_word_embeddings (bool)

classmethod from_hf(hf)[source]#
Parameters:

hf (dict)

Return type:

Qwen3LMConfig

head_dim: int = 128#
hidden_size: int = 2048#
intermediate_size: int = 6144#
max_position_embeddings: int = 32768#
num_attention_heads: int = 16#
num_hidden_layers: int = 28#
num_key_value_heads: int = 8#
rms_norm_eps: float = 1e-06#
rope_theta: float = 1000000.0#
tie_word_embeddings: bool = False#
vocab_size: int = 151936#