nnterp · encyclopedia mixtral · nnterp.families.mixtral

family mixtral

Mixtral

Mistral's pre-norm block with a mixture of 8 SwiGLU experts in place of the MLP: each token goes to 2 of them, whose softmax weights are renormalized to sum to one, and no shared expert.

Mixture of expertsGrouped-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 MixtralForCausalLMlayoutwhat 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] MixtralDecoderLayerlayoutwhat 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 MixtralAttentionlayoutwhat 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 MixtralSparseMoeBlocklayoutwhat 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

Printout

What print(model) shows on mistralai/Mixtral-8x7B-v0.1: the native modules, standard names and every standard value with its layout.

MixtralForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(32000, 4096)
  (layers): ModuleList(
    (0-31): 32 x MixtralDecoderLayer(
      (self_attn): MixtralAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, 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): MixtralSparseMoeBlock(
        (router/gate): MixtralTopKRouter()
        (experts): MixtralExperts(
          (act_fn): SiLUActivation()
        )
        (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
        (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): MixtralRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm): MixtralRMSNorm((4096,), eps=1e-05)
      (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): MixtralRMSNorm((4096,), eps=1e-05)
  (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_size32000
model.intermediate_size14336
architecturesMixtralForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size14336
num_attention_heads32
num_key_value_heads8
vocab_size32000
max_position_embeddings32768
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
num_experts_per_tok2
num_local_experts8

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))          # grouped-query attention, no window
out = h + mlp(post_attention_layernorm(h))       # 2 of 8 experts per token

Mistral's pre-norm block with Llama's names, and a mixture where Mistral has its MLP. Nothing norms a sublayer's output, and the mixture has no shared expert, so mlp_output is routed_output and 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()
    routed = model.layers[1].mlp.routed_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

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

Every block of every checkpoint is a mixture. The experts are intermediate_size wide (14336 on 8x7B, 16384 on 8x22B), and there is no dense width beside it.

The router renormalizes its top 2, in float32

The router (gate, aliased router) gives one logit per expert. It takes a softmax over the 8 in float32, keeps the 2 largest and divides them by their sum, so a token's two expert_weights sum to one; there is no config flag for it. The weights stay float32 whatever the load dtype: in a bfloat16 load router_logits and routed_output are bfloat16 and expert_weights float32.

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

top = logits.float().softmax(-1).topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values / top.values.sum(-1, keepdim=True), w)
assert torch.equal(top.indices, idx)        # slot 0 the larger weight

moe.SCORING is "softmax". With two slots that sum to one, a token's weights are p and 1 - p: the mixture's output is a convex combination of two experts' outputs. router_jitter_noise is 0 on every checkpoint, and the noise, a random factor on the mixture's input, applies only in training.

Ablating an expert

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other weight as it was, so the token keeps one expert scaled by its own weight, below one. mlp_output 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

This holds on the pinned tiny checkpoint for every expert, in float32 and in bfloat16. To give the token's remaining expert the whole weight, write 1 into the other slot as well.

expert_outputs needs the default experts implementation

The 8 experts are one module, experts, with their weights stacked (gate_up_proj, down_proj). expert_outputs, each slot's weighted output, is read inside transformers' grouped_mm (the default) and batched_mm forwards; under experts_implementation="eager" support() reports it missing and names the kwarg. The other mixture values read the same under every implementation, and expert_outputs sums over its slot axis to routed_output.

Attention has no window

Every checkpoint here sets sliding_window to null, so every block attends causally over the whole prompt with transformers' plain causal mask. 8x7B has 32 query heads over 8 key/value heads, 8x22B 48 over 8, so an edit to key/value head j reaches query heads 4j to 4j + 3 on 8x7B and 6j to 6j + 5 on 8x22B. head_dim is 128 on both, and no projection has a bias. The attention interior needs attn_implementation="eager". rope_theta is 10⁶ on every checkpoint.

The tokenizer prepends a BOS, and the chat templates write one

The tokenizers prepend <s> (id 1), so position 0 is the BOS. The Instruct templates begin the text with the string <s> (<s> [INST] ... [/INST] on 8x7B, <s>[INST] ... [/INST] on 8x22B), and a string passed to model.trace is tokenized with special tokens, so a templated prompt starts with two BOS tokens. Tokenize it without them:

text = model.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
batch = model.tokenizer(text, add_special_tokens=False, return_tensors="pt")
with model.trace(dict(batch)):
    ids = model.input_ids.save()            # one <s>

8x22B-Instruct-v0.1 has 32768 tokens, the others 32000; its [INST] and [/INST] are single control tokens.

Target tokens differ by tokenizer

"Paris" encodes to ▁Paris on every checkpoint. " Paris" is ▁Paris too on 8x7B-Instruct and both 8x22B tokenizers, but two tokens, ['▁', '▁Paris'], on Mixtral-8x7B-v0.1 and the pinned tiny checkpoint. Check that the target is one token on the checkpoint you run:

ids = model.tokenizer(" Paris", add_special_tokens=False).input_ids
assert len(ids) == 1, model.tokenizer.convert_ids_to_tokens(ids)

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/mixtral.py

Mixtral (MixtralForCausalLM).

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 (Moe); its router, gate, is aliased router.