nnterp · encyclopedia gpt_neo · nnterp.families.gpt_neo

family gpt_neo

GPT-NeoTinyStories

GPT-2's sequential block and learned position embedding; the attention sits one module down, inside an attn wrapper, computes its own scores without scaling them, and on every other block masks keys outside a local window.

Sliding windowPosition embeddingsLayerNorm

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by attention_layers.

embed_tokenstoken_embeddings
→
layers× 12
→
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 GPTNeoForCausalLMlayoutwhat 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] GPTNeoBlocklayoutwhat 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].mlp GPTNeoMLPlayoutwhat 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
model.layers[i].self_attn GPTNeoSelfAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of self__attn_0, element 0
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of self__attn_0, element 1
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of self__attn_0, element 2
attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked and unscaled, in fp32

inside the forward: the input of nn_functional_softmax_0 inside self__attn_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 self_attn_dropout_0 inside self__attn_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 self__attn_0, element 0

Printout

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

GPTNeoForCausalLM(
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50257, 768)
  (layers): ModuleList(
    (0-11): 12 x GPTNeoBlock(
      (input_layernorm/ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (attn): GPTNeoAttention(
        (attention): GPTNeoSelfAttention(
          (attn_dropout): Dropout(p=0.0, inplace=False)
          (resid_dropout): Dropout(p=0.0, inplace=False)
          (k_proj): Linear(in_features=768, out_features=768, bias=False)
          (v_proj): Linear(in_features=768, out_features=768, bias=False)
          (q_proj): Linear(in_features=768, out_features=768, bias=False)
          (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 and unscaled, in fp32
          (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/ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (mlp): GPTNeoMLP(
        (c_fc): Linear(in_features=768, out_features=3072, bias=True)
        (down_proj/c_proj): Linear(in_features=3072, out_features=768, bias=True)
        (act_fn/act): NewGELUActivation()
        (dropout): Dropout(p=0.0, inplace=False)
        (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
      )
      (self_attn): GPTNeoSelfAttention(
        (attn_dropout): Dropout(p=0.0, inplace=False)
        (resid_dropout): Dropout(p=0.0, inplace=False)
        (k_proj): Linear(in_features=768, out_features=768, bias=False)
        (v_proj): Linear(in_features=768, out_features=768, bias=False)
        (q_proj): Linear(in_features=768, out_features=768, bias=False)
        (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 and unscaled, in fp32
        (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
      )
      (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((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_size50257
model.intermediate_size3072
architecturesGPTNeoForCausalLM
hidden_size768
vocab_size50257
max_position_embeddings2048
tie_word_embeddingsTrue
activation_functiongelu_new
layer_norm_epsilon1e-05
window_size256

The Notes

03 · what to know before you run

The block, in order

h   = x + attn(ln_1(x))      # attn returns attn.attention's output; ln_1 = input_layernorm
out = h + mlp(ln_2(h))       # ln_2 = post_attention_layernorm

transformer.h[i].attn is a GPTNeoAttention wrapper that calls attn.attention, a GPTNeoSelfAttention, and returns its output unchanged. self_attn names the inner module: the projections, the arithmetic and the six interior values are there, and attention_output is its first output, after out_proj. Inside one trace, read it before the wrapper's attn.output. The block returns (hidden_states, attn_weights) on every call, so model.layers[i].output[1] is the pattern and layer_output is the first element. Nothing sits between a sublayer and the stream, and layers[i].input + attention_output + mlp_output == layer_output is exact (difference 0.0 on the pinned checkpoint).

Local and global blocks alternate

config.attention_layers names each block global or local; every released checkpoint with more than one block starts with global and alternates. A local block masks every key window_size or more tokens back (256 on every released checkpoint), so a query attends to itself and the 255 tokens before it. The window is built into the causal mask the module holds (its bias buffer), so outside the window the pattern is zero and the scores hold float32's minimum. A prompt of 256 tokens or fewer gives a local block the same mask as a global one.

window = model.config.window_size
local = model.config.attention_layers.index("local")
prompt = " ".join(["token"] * (window + 4))

with model.trace(prompt):
    pattern = model.layers[local].self_attn.attention_probabilities.save()

n = pattern.shape[-1]
outside = torch.ones(n, n, dtype=torch.bool, device=pattern.device).tril(-window)
assert pattern[..., outside].abs().max() == 0

The scores are not scaled

The module's _attn(query, key, value, mask) computes q @ kᵀ in float32 with no 1 / sqrt(head_dim), applies its own causal (and local) mask with torch.where, adds the model's attention mask (causal, with any padding), takes the softmax, casts it back to the values' dtype and multiplies the values. Its arguments are attention_queries, attention_keys and attention_values (num_heads each, no grouping), the softmax's input is attention_scores and its first return is attention_head_outputs.

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

n = scores.shape[-1]
causal = torch.ones(n, n, dtype=torch.bool, device=scores.device).tril()
again = q.float() @ k.float().transpose(-1, -2)          # no division by sqrt(head_dim)
assert torch.equal(again[..., causal], scores[..., causal])

The scores are large: at block 6 of gpt-neo-125m their standard deviation over the unmasked entries is 10.8 and their largest entry 38 on a 47-token prompt. On a float32 or bfloat16 load the scores above the diagonal are -inf, not float32's minimum: the module's own mask writes float32's minimum there, and the model's attention mask, added next, holds the minimum of the model's dtype, so the sum overflows. attention_scores * 0 is then NaN above the diagonal. On a float16 load the sum stays at float32's minimum. To edit the scores, add to them, or leave the masked entries as they are.

Loading: eager is the default

transformers has eager and flash_attention_2 for GPT-Neo and no sdpa, so a load with no attn_implementation runs eager and serves all six interior values. They are read inside the eager attention, so a flash_attention_2 load reports them unavailable. The EleutherAI configs name no dtype and the TinyStories ones name float32, so a load with no dtype is float32.

Position embeddings are added after embed_tokens

token_embeddings is wte's lookup only. The learned position embedding transformer.wpe, one row per position up to max_position_embeddings (2048), is added to it before block 0, so an edit to token_embeddings leaves the position term in place:

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

assert torch.equal(tokens + positions, x0)

The first position takes most of the global blocks' attention

On gpt-neo-125m, on a 47-token prompt, the global blocks 6, 8 and 10 put 63% to 76% of their pattern (averaged over heads and the later queries) on key 0, and the local blocks between them 23% to 35%. The stream at position 0 has a norm of about 9,100 to 9,500 from block 5's output to block 9's, against medians of about 1,500 at the other positions. The tokenizer is GPT-2's and adds no BOS, so the prompt's first token becomes the sink; leave position 0 out of norms and activation statistics.

The readout: LayerNorm with a bias, tied weights

logits is lm_head.output with nothing applied, and project_on_vocab applied to the last block's layer_output equals logits exactly (difference 0.0 on gpt-neo-125m). lm_head has no bias and its weight is embed_tokens.weight, the same tensor. ln_f's bias adds lm_head.weight @ norm.bias to every position's logits whatever the input: on gpt-neo-125m that vector has a standard deviation of 1.34 logits and its top tokens are the, ,, \n, and, ".

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

The family's checkpoints

gpt-neo-125m is 12 blocks × 768 (12 heads of 64), gpt-neo-1.3B 24 × 2048 and gpt-neo-2.7B 32 × 2560 (16 and 20 heads of 128). Every one has gelu_new, an MLP four times the width (intermediate_size is None, which means 4 * hidden_size) and GPT-2's 50257-token vocabulary. The TinyStories models have 1 to 8 blocks, 64 to 1024 wide, always 16 heads, so head_dim runs from 4 (TinyStories-1M) to 64; TinyStories-1Layer-21M's one block is global.

From the family module

04 · nnterp/families/gpt_neo.py

GPT-Neo (GPTNeoForCausalLM).

transformer.{wte, wpe, drop, h[i].{ln_1, attn.attention, ln_2, mlp}, ln_f} and lm_head. A sequential block that returns (hidden_states, attn_weights). attn is a GPTNeoAttention wrapper that forwards to attn.attention, the GPTNeoSelfAttention holding the projections, and returns its output as is; self_attn is bound to that inner module, so its input is ln_1's output and its first output is what the block adds. The inner module does its own arithmetic in _attn: an fp32 q @ k^T with no 1/sqrt(head_dim) scaling, a causal mask by torch.where (a sliding window of window_size on the local layers of attention_layers, which alternate with global ones), the padding mask added, softmax, a cast to the values' dtype, attn_dropout, then @ v. The model knows only eager and flash_attention_2; the interior is read around the _attn call and needs eager. wpe and drop have no standard name.