nnterp · encyclopedia qwen2_moe · nnterp.families.qwen2_moe

family qwen2_moe

Qwen1.5-MoEQwen2-MoE

Qwen2's attention, with biased q, k and v, and a mixture of experts on every block: the router's top-k softmax weights are not renormalized, and a shared expert is added through a sigmoid gate of its own.

Biased q, k, vMixture of expertsUnnormalized routing

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model Qwen2MoeForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] Qwen2MoeDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn Qwen2MoeAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp Qwen2MoeSparseMoeBlocklayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output
shared_expert_output Residual [batch seq hidden_size] The shared expert's output

inside the forward: the output of shared_expert_output_1

Printout

What print(model) shows on Qwen/Qwen1.5-MoE-A2.7B: the native modules, standard names and every standard value with its layout.

Qwen2MoeForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=151936, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(151936, 2048)
  (layers): ModuleList(
    (0-23): 24 x Qwen2MoeDecoderLayer(
      (self_attn): Qwen2MoeAttention(
        (q_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (k_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (v_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Qwen2MoeSparseMoeBlock(
        (router/gate): Qwen2MoeTopKRouter()
        (experts): Qwen2MoeExperts(
          (act_fn): SiLUActivation()
        )
        (shared_experts/shared_expert): Qwen2MoeMLP(
          (gate_proj): Linear(in_features=2048, out_features=5632, bias=False)
          (up_proj): Linear(in_features=2048, out_features=5632, bias=False)
          (down_proj): Linear(in_features=5632, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
        )
        (shared_expert_gate): Linear(in_features=2048, out_features=1, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers24
model.hidden_size2048
model.num_heads16
model.num_kv_heads16
model.head_dim128
model.qk_head_dim128
model.vocab_size151936
model.intermediate_size5632
architecturesQwen2MoeForCausalLM
num_hidden_layers24
hidden_size2048
intermediate_size5632
num_attention_heads16
num_key_value_heads16
vocab_size151936
max_position_embeddings8192
sliding_window0
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
moe_intermediate_size1408
num_experts60
num_experts_per_tok4
norm_topk_probFalse
use_sliding_windowFalse
max_window_layers21
decoder_sparse_step1
shared_expert_intermediate_size5632

The Notes

03 · what to know before you run

The block, in order

h      = x + self_attn(input_layernorm(x))       # q, k, v projections biased
y      = post_attention_layernorm(h)
shared = shared_expert(y)                        # runs first
routed = experts(y, router(y))                   # 4 of 60 experts per token (1.5-MoE)
out    = h + routed + sigmoid(shared_expert_gate(y)) * shared

Qwen2's pre-norm block with Llama's names. The mixture's output is the routed sum plus the gated shared expert, and no norm follows either sublayer. The identity is the plain sum, and inside the mixture routed_output + shared_expert_output is mlp_output:

moe = model.layers[1].mlp
with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    mlp = moe.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(routed + shared, mlp)
torch.testing.assert_close(x + attn + mlp, out)

The shared expert is gated by a sigmoid

shared_experts (native shared_expert) is a SwiGLU MLP shared_expert_intermediate_size wide (5632 on 1.5-MoE, 20480 on 57B-A14B) that every token runs through. shared_expert_gate is a linear map from the stream to one number per token; the mixture multiplies the shared expert's output by its sigmoid. shared_expert_output is that product, what the mixture adds; moe.shared_experts.output is the ungated output:

moe = model.layers[1].mlp
with model.trace(prompt):
    ungated = moe.shared_experts.output.save()
    gate = moe.shared_expert_gate.output.save()      # [batch * seq, 1], before the sigmoid
    shared = moe.shared_expert_output.save()

torch.testing.assert_close(torch.sigmoid(gate).view(*shared.shape[:-1], 1) * ungated, shared)

Zeroing shared_expert_output removes the shared expert and leaves mlp_output equal to routed_output; zeroing shared_experts.output does the same through the product.

Read order: the shared expert runs before the router

The mixture's forward calls shared_expert first, then the router and the experts, and computes the gated product last. In one trace moe.shared_experts.output comes before router_logits, and shared_expert_output after routed_output; the other orders raise OutOfOrderError:

moe = model.layers[1].mlp
with model.trace(prompt):
    ungated = moe.shared_experts.output.save()
    logits = moe.router_logits.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()

The routed weights are the softmax's top-k, not renormalized

The router (gate, aliased router) takes a softmax over every expert's logit in float32 and keeps the largest num_experts_per_tok; norm_topk_prob is false on every checkpoint here, so expert_weights are those probabilities as they are, cast back to the model's dtype, and a token's weights sum to less than one:

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()       # [batch, seq, num_experts], before the softmax
    w = moe.expert_weights.save()           # [batch, seq, top_k]
    idx = moe.expert_indices.save()

top = logits.float().softmax(-1).topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values.to(w.dtype), w)    # the softmax's own entries
assert torch.equal(top.indices, idx)                     # slots in descending weight

1.5-MoE has 60 experts with 4 per token, each moe_intermediate_size 1408 wide; 57B-A14B has 64 with 8 per token, 2560 wide. The shared expert's gate does not depend on the routed weights, so the balance between the routed sum and the shared expert moves with how much softmax mass the router puts on the experts it chose.

Ablating an expert

Zeroing expert e's weight where expert_indices == e removes that slot's term. mlp_output changes on exactly the tokens that chose e, and an expert no token chose has no effect:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The check holds on the pinned tiny checkpoint for every expert, in float32 and in bfloat16. expert_outputs is read inside transformers' grouped_mm (the default) and batched_mm experts forwards; under experts_implementation="eager" support() reports it missing and names the kwarg.

Every block is a mixture

A block's mlp is a mixture where its index is not in mlp_only_layers and (i + 1) % decoder_sparse_step == 0. No checkpoint here sets mlp_only_layers and every one has decoder_sparse_step 1, so all 24 blocks of 1.5-MoE and 28 of 57B-A14B are mixtures. The config's intermediate_size (5632, 18944) is the width of a dense MLP no block has; model.layers[i].mlp.intermediate_size reads the experts' moe_intermediate_size.

Attention

q_proj, k_proj and v_proj add a bias (qkv_bias is true by default and no checkpoint sets it); o_proj does not. The bias is added before the rotary embedding, so it turns with the position. 1.5-MoE has 16 query heads and 16 key/value heads, no grouping; 57B-A14B has 28 over 4, so an edit to key/value head j reaches query heads 7j to 7j + 6. head_dim is 128 on both. Every checkpoint sets use_sliding_window false, so every entry of layer_types is full_attention and no block has a window, whatever sliding_window says. The attention interior needs attn_implementation="eager".

Tokenizer and chat template

The tokenizer prepends nothing (bos_token is None), so position 0 is the text's first token. Every checkpoint ships a ChatML template that writes a system turn when given none: You are a helpful assistant on 1.5-MoE-A2.7B, the same with a full stop on the Chat and Qwen2 models.

The readout

model.logits equals model.lm_head.output, and project_on_vocab on the last block's layer_output equals logits. lm_head and embed_tokens are separate weights on every checkpoint here.

From the family module

04 · nnterp/families/qwen2_moe.py

Qwen2-MoE (Qwen2MoeForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward. The MLP is a sparse mixture of experts with a shared expert.

A config with mlp_only_layers set gives those blocks a dense Qwen2MoeMLP that no envoy is keyed to: layers[i].mlp is a plain Envoy there and support() reports "no mlp module on this block". No released checkpoint sets it.