nnterp · encyclopedia mistral · nnterp.families.mistral

family mistral

MistralMistral NemoMistral Small

Llama's pre-norm block with grouped-query attention; on Mistral-7B-v0.1 every block attends over a sliding window of the last 4096 positions, and from Nemo on the attention is narrower than the stream.

Sliding windowGrouped-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 MistralForCausalLMlayoutwhat 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] MistralDecoderLayerlayoutwhat 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 MistralAttentionlayoutwhat 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 MistralMLPlayoutwhat 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

Printout

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

MistralForCausalLM(
  (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 MistralDecoderLayer(
      (self_attn): MistralAttention(
        (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): MistralMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): MistralRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm): MistralRMSNorm((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): MistralRMSNorm((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
architecturesMistralForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size14336
num_attention_heads32
num_key_value_heads8
head_dim128
vocab_size32000
max_position_embeddings32768
sliding_window4096
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))
out = h + mlp(post_attention_layernorm(h))

Llama's block: two RMSNorms per block, each before its sublayer, none after, and a SwiGLU MLP (gate_proj, up_proj, down_proj, SiLU). post_attention_layernorm is named for where it sits on the stream, after the attention's add: it is the MLP's input norm.

The contributions are the modules' outputs

attention_output is the attention module's output and mlp_output the MLP's, added to the stream with nothing in between, so the identity holds exactly (float32, pinned checkpoint):

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)

A sliding window on Mistral-7B-v0.1

config.sliding_window is 4096 on Mistral-7B-v0.1 and BakLLaVA and None on v0.3, Nemo, Small 3.1 and 3.2, Pixtral-12B and llava-v1.6-mistral. The config has no layer_types: where the window is set, every block takes it. A query attends to its own position and the 4095 before it, so on a prompt shorter than the window the mask is the plain causal mask. The window is in the mask the attention receives, so under attn_implementation="eager" attention_probabilities is exactly zero beyond it (checked on the pinned checkpoint loaded with sliding_window=3: every entry three or more positions back is 0.0).

model.config.get_text_config().sliding_window   # 4096 on v0.1, None on v0.3

Grouped-query attention, narrower than the stream from Nemo on

Every checkpoint listed has 32 query heads over 8 key/value heads of width 128. attention_keys and attention_values are read before repeat_kv, [batch, 8, seq, 128], and query head h reads key/value head h // 4: an edit to key/value head j reaches query heads 4j to 4j + 3. On the 7B checkpoints 32 × 128 is the stream's width, 4096. On Nemo, Small 3.1 and 3.2 and Pixtral-12B the stream is 5120 wide and the heads still 32 × 128: q_proj maps 5120 to 4096 and o_proj 4096 back to 5120, so attention_head_outputs flattened is 4096 wide, not hidden_size.

Target tokens differ by tokenizer

Mistral-7B v0.1 and v0.3 use a SentencePiece vocabulary (32000 and 32768 tokens): " Paris" and "Paris" both encode to ▁Paris (5465 on v0.1, 6233 on v0.3). Nemo, Small 3.1 and Pixtral-12B use the 131072-token Tekken vocabulary, where they are two tokens: " Paris" is ĠParis (6993) and "Paris" is Paris (42572). The pinned tiny checkpoint's tokenizer encodes " Paris" to two tokens, ['▁', '▁Paris']. Decode the target ids 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 and the embeddings

model.logits is model.lm_head.output: no cap and no scale, so project_on_vocab is lm_head(norm(hidden)). token_embeddings is layers[0].input on a text prompt (checked on the pinned checkpoint). No checkpoint listed ties lm_head to embed_tokens.

What this family module covers

The module serves every checkpoint whose config resolves to model_type mistral: the 7B base and instruct models v0.1 to v0.3, Mistral Nemo, Mistral Small 24B, and the text model of four vision-language wrappers (Mistral Small 3.1 and 3.2, Pixtral-12B, llava-v1.6-mistral, BakLLaVA). mistralai/Ministral-8B-Instruct-2410's config.json says mistral, but AutoConfig loads a mistral config that has layer_types as a ministral config, so it is the ministral family. Mixtral is the mixtral family, and Ministral 3's text model ministral3.

From the family module

04 · nnterp/families/mistral.py

Mistral (MistralForCausalLM).

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.

Two towers sit at model.vision_tower, which is vision, with the projector model.multi_modal_projector as projector: Pixtral (a PixtralVision, every image's patches packed in one row, no final norm) on the mistral3 wrapper (Mistral Small 3.1 / 3.2) and on Pixtral-12B's llava wrapper; CLIP on LLaVA-NeXT (llava-v1.6-mistral) and BakLLaVA (llava). The towers' inner names differ, so both sets are keyed. vision.image_features is read at the scatter (ImageScatter), so on LLaVA-NeXT it is the unpadded projector output with its newline tokens, as the text model receives it. Pixtral-12B reads vision_feature_layer -1 with strategy "full", so its projector.input is vision.tower_output; BakLLaVA reads block -2 without the CLS token, as Llava 1.5.