nnterp · encyclopedia mpt · nnterp.families.mpt

family mpt

MPT

The block adds the attention to the stream, but the MLP adds the residual inside its own module, so mlp_output is the tensor before that add; no projection or norm has a bias.

ALiBiLayerNorm

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

model MptForCausalLMlayoutwhat 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] MptBlocklayoutwhat 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 MptAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the output of query_states_0
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the output of key_states_0
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the output of value_states_0
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
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
attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of torch_matmul_1
model.layers[i].mlp MptMLPlayoutwhat 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 tensor before the residual is added inside

inside the forward: the output of F_dropout_0

Printout

What print(model) shows on ybelkada/mpt-7b-bf16-sharded: the native modules, standard names and every standard value with its layout.

MptForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=50432, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50432, 4096)
  (layers): ModuleList(
    (0-31): 32 x MptBlock(
      (input_layernorm/norm_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=False)
      (self_attn/attn): MptAttention(
        (Wqkv): Linear(in_features=4096, out_features=12288, bias=False)
        (o_proj/out_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_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (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
      )
      (post_attention_layernorm/norm_2): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=False)
      (mlp/ffn): MptMLP(
        (up_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (act_fn/act): GELU(approximate='none')
        (down_proj): Linear(in_features=16384, out_features=4096, bias=False)
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the tensor before the residual is added inside
        (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
      )
      (resid_attn_dropout): Dropout(p=0, inplace=False)
      (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=False)
  (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_heads32
model.head_dim128
model.qk_head_dim128
model.vocab_size50432
model.intermediate_size16384
architecturesMptForCausalLM
vocab_size50432
tie_word_embeddingsTrue
dtypebfloat16
layer_norm_epsilon1e-05
d_model4096
n_heads32
n_layers32
max_seq_len2048

The Notes

03 · what to know before you run

The MLP adds the residual inside the module

h   = x + self_attn(norm_1(x))                 # the block adds the attention
out = ffn(norm_2(h), residual=h)               # returns h + mlp_output

The attention returns its out_proj output and the block adds it, so attention_output is the module's output. The MLP takes the stream as a second argument and ends in output + residual, so mlp.output is the block's output, equal to layer_output, and mlp_output is the tensor before that add: down_proj's output after its dropout. The sum is exact in every dtype.

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()
with model.trace(prompt):
    raw = model.layers[1].mlp.output.save()

assert torch.equal(raw, x + attn + mlp)       # the module returns the stream
assert torch.equal(x + attn + mlp, out)

Zeroing mlp.output zeroes the stream that leaves the block, not the MLP's contribution; ablate, steer or patch mlp_output. mlp_output is an operation inside the module, so in one trace read it before mlp.output: the other order raises OutOfOrderError naming ffn.source.F_dropout_0.output.

Queries, keys and values are thirds of one projection

Wqkv maps hidden_size to 3 * hidden_size with no bias, and chunk(3) cuts its output into queries, keys and values, each then split into heads: the first third is every head's query, the second every key, the third every value. No grouping: keys and values have num_heads heads.

attn = model.layers[1].self_attn
with model.trace(prompt):
    h = model.layers[1].input_layernorm.output.save()
    q = attn.attention_queries.save()

fused = h @ attn.Wqkv.weight.T                                   # [batch, seq, 3 * hidden_size]
q_part, k_part, v_part = fused.chunk(3, dim=-1)
heads = q_part.view(*q_part.shape[:2], model.num_heads, model.head_dim).transpose(1, 2)
torch.testing.assert_close(heads, q)

The three are views of the outputs of one chunk, and torch refuses an in-place edit of them (RuntimeError: Output 0 of Transpose is a view and is being modified inplace). Assign instead:

attn = model.layers[1].self_attn
with model.trace(prompt):
    attn.attention_queries = attn.attention_queries * 0
    logits = model.logits.save()

ALiBi biases replace rotary

There is no rotary embedding and no position embedding: queries and keys carry no position, and the same token at two positions gives the same attention_queries at block 0. Position enters as a bias on the scores: attention_scores is q @ kᵀ / sqrt(head_dim) plus slope_h * (j - (seq - 1)) for key position j, then the mask. The bias is 0 on the last key, negative before it, and the same on every query row; after the softmax it is a penalty of slope_h per token of distance. Head h has the slope 2 ** (-8 * (h + 1) / heads): on MPT-7B, with 32 heads, 0.84 for head 0 down to 0.0039 for head 31. transformers builds the bias with alibi_bias_max 8 whatever the config says. On a float32 load the scores recompute to rounding:

from transformers.models.mpt.modeling_mpt import build_mpt_alibi_tensor

attn = model.layers[1].self_attn
with model.trace(prompt):
    q = attn.attention_queries.save()
    k = attn.attention_keys.save()
    scores = attn.attention_scores.save()

seq = q.shape[2]
alibi = build_mpt_alibi_tensor(model.num_heads, seq)                  # [num_heads, 1, seq]
again = q.float() @ k.float().transpose(-1, -2) / model.head_dim ** 0.5 + alibi
causal = torch.ones(seq, seq, dtype=torch.bool).tril()
torch.testing.assert_close(again[..., causal], scores[..., causal])

The attention needs no flag, and its scores are float32

transformers has one attention for MPT, its own arithmetic, so a load without attn_implementation serves all six interior values: support() reports none missing. The scores are read where the softmax takes them, after an upcast, so attention_scores is float32 whatever the model's dtype; the masked entries hold the model dtype's minimum (-3.39e38 under bfloat16). The pattern is cast back to the values' dtype. The head outputs are the second matmul, heads first in the forward, served [batch, seq, num_heads, head_dim].

What transformers' MptForCausalLM reads from the config

The class is transformers' own; the checkpoints' auto_map code is not run. It takes the sizes, attn_config.softmax_scale and attn_config.clip_qkv, and builds every Linear without a bias and every LayerNorm without one. The MLP is 4 * hidden_size wide; every listed checkpoint sets expansion_ratio to 4. ethzanalytics/mpt-7b-storywriter-sharded sets clip_qkv to the integer 6, which transformers' MptAttentionConfig refuses, so its config does not load.

The readout

lm_head is wte's matrix (tied) with no bias, and logits equals lm_head.output; norm_f has no bias, and project_on_vocab on the last block's layer_output equals logits.

From the family module

04 · nnterp/families/mpt.py

MPT (MptForCausalLM).

transformer.{wte, blocks[i].{norm_1, attn, norm_2, ffn, resid_attn_dropout}, norm_f} and lm_head. The block adds the attention's output itself, but the MLP takes the residual and adds it inside, so its contribution is the tensor before that add (its dropout's output). The attention does its own arithmetic whatever attn_implementation says, so the pattern is its dropout's output and needs no eager load.