nnterp · encyclopedia opt · nnterp.families.opt

family opt

OPT

A LayerNorm block with a learned position embedding offset by 2 and no MLP module: fc1, the activation and fc2 sit on the block itself, and fc2.output is what the MLP path adds to the stream.

No MLP modulePosition embeddingsLayerNormBiased q, k, v

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× 12
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + fc2.output.view_as(layer_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 OPTForCausalLMlayoutwhat 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] OPTDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output scalar The normed stream entering the attention

the child module input_layernorm's output
layer_mid FlatResidual [tokens hidden_size] The residual stream after the attention add, flattened to [batch * seq, hidden]

the child module final_layer_norm's input
mlp_norm_output scalar The normed stream entering the feed-forward path, flattened to [batch * seq, hidden]

the child module final_layer_norm's output
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 OPTAttentionlayoutwhat 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

Printout

What print(model) shows on facebook/opt-125m: the native modules, standard names and every standard value with its layout.

OPTForCausalLM(
  (lm_head): Linear(in_features=768, out_features=50272, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50272, 768, padding_idx=1)
  (layers): ModuleList(
    (0-11): 12 x OPTDecoderLayer(
      (self_attn): OPTAttention(
        (k_proj): Linear(in_features=768, out_features=768, bias=True)
        (v_proj): Linear(in_features=768, out_features=768, bias=True)
        (q_proj): Linear(in_features=768, out_features=768, bias=True)
        (o_proj/out_proj): Linear(in_features=768, out_features=768, 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
      )
      (activation_fn): ReLU()
      (input_layernorm/self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (fc1): Linear(in_features=768, out_features=3072, bias=True)
      (fc2): Linear(in_features=3072, out_features=768, bias=True)
      (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (attention_norm_output): The normed stream entering the attention
      (layer_mid) -> FlatResidual [tokens hidden_size]: The residual stream after the attention add, flattened to [batch * seq, hidden]
      (mlp_norm_output): The normed stream entering the feed-forward path, flattened to [batch * seq, hidden]
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (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((768,), 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_layers12
model.hidden_size768
model.num_heads12
model.num_kv_heads12
model.head_dim64
model.qk_head_dim64
model.vocab_size50272
model.intermediate_size3072
architecturesOPTForCausalLM
num_hidden_layers12
hidden_size768
num_attention_heads12
vocab_size50272
max_position_embeddings2048
tie_word_embeddingsTrue
dtypefloat16
activation_functionrelu
ffn_dim3072
do_layer_norm_beforeTrue
word_embed_proj_dim768

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(self_attn_layer_norm(x))           # self_attn_layer_norm = input_layernorm
f   = h.reshape(-1, hidden)                            # [batch * seq, hidden_size]
out = (f + fc2(relu(fc1(final_layer_norm(f))))).view_as(x)

The block returns a bare tensor. attention_output is the attention module's output after out_proj. The block's own final_layer_norm is the MLP path's input norm; it keeps its native name, since the decoder's last norm has the same name and is model.norm. Every norm is a LayerNorm with a weight and a bias, and every projection, q_proj, k_proj, v_proj, out_proj, fc1 and fc2, has a bias.

The MLP is fc1 and fc2 on the block

No block has an mlp module, so mlp, mlp_output and their kin do not exist on OPT and support() lists no mlp.* key. What the MLP path adds is layers[i].fc2.output, and on every pre-norm size layers[i].fc1.input is final_layer_norm's output. The block flattens the stream before that path, so fc1.output, activation_fn.output (the neurons) and fc2.output, and on a pre-norm size final_layer_norm.output, are [batch * seq, ...], row b * seq + t for prompt b, position t. Unflatten them to read or edit a position:

with model.trace(prompts):
    batch, seq = model.input_size
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    neurons = model.layers[1].activation_fn.output.save()   # [batch * seq, ffn_dim]
    fc2 = model.layers[1].fc2.output
    mlp_out = fc2.unflatten(0, (batch, seq)).save()
    out = model.layers[1].layer_output.save()

torch.equal(x + attn + mlp_out, out)                        # True

The unflattened tensor is a view, so an in-place edit lands: model.layers[1].fc2.output.unflatten(0, (batch, seq))[:, -1] = 0 removes the MLP path's term at the last position, and that position's layer_output is then input + attention_output. intermediate_size is the config's ffn_dim, four times the width on every released size, and the activation is ReLU.

The queries arrive scaled

OPT multiplies the projected queries by head_dim ** -0.5 and calls the attention with a scale of 1, so attention_queries is q_proj's output times that scale, and attention_scores is attention_queries @ attention_keys.transpose(-1, -2) plus the mask. A query-key product computed from q_proj.output is head_dim ** 0.5 times the score. The interior values need attn_implementation="eager"; the default load runs sdpa. Every head has its own keys and values.

Learned positions, offset by 2

model.decoder.embed_positions is an embedding of max_position_embeddings + 2 rows (2050), and position t reads row t + 2. Its output is added to the token embeddings before block 0:

with model.trace(prompt):
    tokens = model.token_embeddings.save()
    positions = model.model.decoder.embed_positions.output.save()
    x0 = model.layers[0].input.save()

W = model.model.decoder.embed_positions.weight
torch.equal(tokens + positions, x0)                         # True
torch.equal(positions[0], W[2 : 2 + x0.shape[1]])           # True

Positions follow the attention mask, so a left-padded prompt reads the same rows as the prompt alone, and its logits match (largest difference 7e-8 on the pinned checkpoint). A prompt longer than 2048 tokens runs past the table.

opt-350m is post-norm and projects the embedding

Of the released sizes only opt-350m sets do_layer_norm_before to false. Each of its norms follows the add: the attention reads the raw stream, self_attn_layer_norm normalizes input + attention_output, and layer_output is the block's final_layer_norm output. The identity does not hold there: on opt-350m, input + attention_output + fc2.output misses layer_output by 4 to 140 in its largest coordinate, depending on the block, on a stream whose norm is 32. The decoder has no final norm, so model.norm does not exist and project_on_vocab raises AttributeError. The embedding is 512 wide: project_in maps token_embeddings to the stream's 1024 before the positions are added, and project_out maps the stream back to 512 for the tied lm_head. The logit lens there is lm_head(project_out(x)), on opt-350m:

with model.trace(prompt):
    last = model.layers[-1].layer_output.save()
    logits = model.logits.save()

project_out = model.model.decoder.project_out._module
lm_head = model.lm_head._module
torch.equal(lm_head(project_out(last)), logits)             # True

The first position is a sink in one coordinate

The tokenizer prepends </s> (id 2) to every prompt. On opt-125m the stream at that position has a norm of 120 after block 0 and 311 to 379 from block 2 to block 10, nearly all of it in coordinate 422, against medians of 2 to 8 at the other positions. From block 1 on, 63% to 85% of the attention pattern, averaged over heads and later queries, lands on key 0. Block 11 brings that position's norm back to 9. Leave position 0 out of norms, steering scales and activation statistics.

The readout: LayerNorm with a bias, tied weights

logits is lm_head.output with nothing applied, and project_on_vocab is lm_head(final_layer_norm(x)), so a logit lens at the last block equals logits on every pre-norm size. lm_head has no bias and its weight is embed_tokens.weight, the same tensor, so an edit to the embedding matrix is an edit to the unembedding.

The family's checkpoints

Blocks × width: 12 × 768 (opt-125m), 24 × 1024 (opt-350m), 24 × 2048 (opt-1.3b), 32 × 2560 (opt-2.7b, head_dim 80), 32 × 4096 (opt-6.7b), 40 × 5120 (opt-13b), 48 × 7168 (opt-30b), 64 × 9216 (opt-66b). head_dim is 64 up to 1.3b and 128 from 6.7b. Every size has ReLU, an MLP four times the width, 2048 positions and the same 50272-token vocabulary, and stores float16 weights.

From the family module

04 · nnterp/families/opt.py

OPT (OPTForCausalLM).

model.decoder.{embed_tokens, embed_positions, layers[i].{self_attn, self_attn_layer_norm, fc1, activation_fn, fc2, final_layer_norm}, final_layer_norm} and lm_head. There is no MLP module: fc1 and fc2 sit on the block, so mlp and mlp_output do not exist here and StandardizedTransformer.support says so. The block's own final_layer_norm is the pre-MLP norm; it keeps its native name, since a single-component alias would also bind on the decoder's final norm.

The MLP path runs on the flattened [batch * seq, ...] stream, so input + attention_output + fc2.output.view_as(layer_output) == layer_output. The queries are scaled by head_dim**-0.5 at q_proj's output, so attention_queries are pre-scaled. facebook/opt-350m sets do_layer_norm_before false (norms after the adds, so that identity does not hold), has no final norm (no norm; project_on_vocab raises AttributeError) and wraps the stack in project_in / project_out (512 <-> 1024): token_embeddings is 512 wide and layers[0].input 1024.