nnterp · encyclopedia phimoe · nnterp.families.phimoe

family phimoe

Phi-3.5-MoE

Llama's tree with LayerNorms and biased projections, and on every block a mixture of 16 experts whose router, sparsemixer, picks two experts in turn and gives each a weight of exactly 1 unless another logit is within 2% of that expert's, so a token's two weights sum to up to 2.

Mixture of expertsBiased q, k, vLayerNormGrouped-query attention

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× 32
→
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 PhimoeForCausalLMlayoutwhat 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] PhimoeDecoderLayerlayoutwhat 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 PhimoeAttentionlayoutwhat 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 PhimoeSparseMoeBlocklayoutwhat 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 forward_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

Printout

What print(model) shows on microsoft/Phi-3.5-MoE-instruct: the native modules, standard names and every standard value with its layout.

PhimoeForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=32064, bias=True)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(32064, 4096)
  (layers): ModuleList(
    (0-31): 32 x PhimoeDecoderLayer(
      (self_attn): PhimoeAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=True)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=True)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=True)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=True)
        (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): PhimoeSparseMoeBlock(
        (router): PhimoeTopKRouter(in_features=4096, out_features=16, bias=False)
        (experts): PhimoeExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
      (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): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
  (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_layers32
model.hidden_size4096
model.num_heads32
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size32064
model.intermediate_size6400
architecturesPhiMoEForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size6400
num_attention_heads32
num_key_value_heads8
vocab_size32064
max_position_embeddings131072
sliding_window131072
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
num_experts_per_tok2
num_local_experts16
attention_biasTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))       # q, k, v, o projections with biases
out = h + mlp(post_attention_layernorm(h))    # 2 of 16 experts per token

Llama's names and Llama's order, with two differences inside: the block's two norms and the final norm are nn.LayerNorms with a weight and a bias (their eps comes from rms_norm_eps), and on every public checkpoint attention_bias puts a bias on q_proj, k_proj, v_proj and o_proj. Nothing norms a sublayer's output, so attention_output is o_proj's output and mlp_output is the mixture's. There is no shared expert: routed_output is mlp_output. The identity is the plain sum:

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

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

The router picks two experts one after the other

mlp.router is an nn.Linear subclass, 16 outputs, no bias; router_logits is its projection, read before the scoring. SCORING is "sparsemixer", and in eval it runs two slots in turn:

  • slot 0 takes the expert with the largest logit m. Every expert whose logit s is within the jitter threshold, (m - s) / max(|s|, m) <= 2 * router_jitter_noise (0.01 on every checkpoint), stays in a softmax; the others are masked out. The weight is the chosen expert's share of that softmax;
  • slot 1 removes slot 0's expert and does the same over the remaining 15.

So expert_indices are the top two logits, and each expert_weights entry is between 0 and 1, exactly 1 when no other logit is within the threshold of that slot's maximum. The two weights are not renormalized together: a token's weights sum to up to 2, and the mixture scales each expert's output by its own. On the pinned checkpoint every weight is 1.0. The recomputation, checked there:

moe = model.layers[1].mlp
eps = model.config.router_jitter_noise
with model.trace(prompt):
    logits = moe.router_logits.save()       # [batch, seq, num_experts]
    w = moe.expert_weights.save()           # [batch, seq, 2]
    idx = moe.expert_indices.save()

def slot(scores):                           # one slot of sparsemixer, in eval
    top = scores.argmax(-1, keepdim=True)
    m = scores.gather(-1, top)
    near = (m - logits) / logits.abs().clamp(min=m) <= 2 * eps
    return scores.masked_fill(~near, float("-inf")).softmax(-1).gather(-1, top), top

w0, i0 = slot(logits)
w1, i1 = slot(logits.scatter(-1, i0, float("-inf")))
assert torch.equal(torch.cat([i0, i1], -1), idx)
torch.testing.assert_close(torch.cat([w0, w1], -1), w)

The weight a near-tie takes is dropped, not passed to the other expert. Written at router_logits, a last token whose logits are 1.0 on expert 2, 0.99 on expert 5 and 0.5 or less elsewhere routes to [2, 5] with weights [0.5025, 1.0]: slot 0 shares its softmax with expert 5, which slot 1 then takes at full weight. Under model.train() the routing is random. The mixture and then its router each multiply the input, in place, by uniform noise in [0.99, 1.01] (input_jitter_noise); each slot samples its expert with Gumbel noise; and a slot that samples an expert other than its largest logit has its weight cut to a third, three times in four. Trace in eval.

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the other slot's weight as it was. mlp_output then changes on exactly the tokens that chose e:

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

A slot whose weight is 1 contributes its expert's whole output, so removing it removes a full-weight term. expert_outputs (each slot's weighted output) is served under the default experts_implementation, and its sum over the two slots is routed_output.

The rotary is longrope with one scale at every length

Phi-3.5-MoE sets rope_scaling type longrope over an original_max_position_embeddings of 4096, with short_mscale and long_mscale both 1.243: cos and sin are multiplied by 1.243 at every position and every length, so a rotated query's norm is 1.243 times the projection's. The transformers rotary for this class computes its frequencies from short_factor at every length; long_factor is not used, also past 4096 tokens. Rotary covers whole 128-dimensional heads. The sliding_window of 131072 equals the context and never masks. Phi-mini-MoE and Phi-tiny-MoE have no rotary scaling, a 4096-token context and a sliding_window of 2047, which masks every position more than 2046 tokens back.

Load with eager for the attention interior only

The attention interior needs attn_implementation="eager"; there is no softcap, sink or window short of the context on Phi-3.5-MoE, so the default sdpa load computes the same function. 32 query heads share 8 key/value heads of 128 dimensions: attention_keys and attention_values are served before repeat_kv, so an edit to key/value head j reaches query heads 4j to 4j + 3.

The readout has two biases

logits equals lm_head.output, and project_on_vocab applied to the last block's layer_output equals logits exactly, the final LayerNorm's bias and lm_head's included. A logit lens therefore adds the same vector at every block, whatever the input:

bias_logits = model.lm_head.weight @ model.norm.bias + model.lm_head.bias

embed_tokens and lm_head are separate matrices. The tokenizer is Phi-3's: 32011 tokens in 32064 rows, <s> (id 1) as bos_token but not prepended, so position 0 is the prompt's first token. ' Paris' splits as ['▁', '▁Paris']; encode "Paris" (id 3681) for one token.

The checkpoints

The sizes, from the configs, each with 16 experts and top 2 on every one of 32 blocks, hidden 4096:

  • Phi-3.5-MoE-instruct: 32 heads × 128 over 8, experts 6400 wide, a 131072-token context
  • Phi-mini-MoE-instruct: 32 heads × 128 over 8, experts 960 wide, a 4096-token context
  • Phi-tiny-MoE-instruct: 16 heads × 128 over 4, experts 448 wide, a 4096-token context

Phi-mini-MoE and Phi-tiny-MoE ship SlimMoE modeling code (configuration_slimmoe, auto_map) beside a model_type of phimoe; without trust_remote_code transformers builds its own PhimoeForCausalLM from the config.

From the family module

04 · nnterp/families/phimoe.py

Phi-3.5-MoE (PhimoeForCausalLM).

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 norms are nn.LayerNorm. Every block's MLP is a sparse mixture of experts that returns the routed hidden states as a bare tensor; the experts are intermediate_size wide.