nnterp · encyclopedia gpt_neox · nnterp.families.gpt_neox

family gpt_neox

GPT-NeoX

Pythia's block is parallel: attention and MLP each read their own LayerNorm of the block input, and the block adds both at once, so a block's MLP never sees that block's attention.

Parallel blockPartial rotary

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× 6
→
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 GPTNeoXForCausalLMlayoutwhat 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] GPTNeoXLayerlayoutwhat 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 scalar 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 GPTNeoXAttentionlayoutwhat 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 GPTNeoXMLPlayoutwhat 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 EleutherAI/pythia-70m-deduped: the native modules, standard names and every standard value with its layout.

GPTNeoXForCausalLM(
  (lm_head): Linear(in_features=512, out_features=50304, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50304, 512)
  (layers): ModuleList(
    (0-5): 6 x GPTNeoXLayer(
      (input_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_dropout): Dropout(p=0.0, inplace=False)
      (post_mlp_dropout): Dropout(p=0.0, inplace=False)
      (self_attn/attention): GPTNeoXAttention(
        (query_key_value): Linear(in_features=512, out_features=1536, bias=True)
        (o_proj/dense): Linear(in_features=512, out_features=512, 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
      )
      (mlp): GPTNeoXMLP(
        (dense_h_to_4h): Linear(in_features=512, out_features=2048, bias=True)
        (down_proj/dense_4h_to_h): Linear(in_features=2048, out_features=512, bias=True)
        (act_fn/act): GELUActivation()
        (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
      )
      (layer_mid): The residual stream after the attention add
      (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
      (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((512,), 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_layers6
model.hidden_size512
model.num_heads8
model.num_kv_heads8
model.head_dim64
model.qk_head_dim64
model.vocab_size50304
model.intermediate_size2048
architecturesGPTNeoXForCausalLM
num_hidden_layers6
hidden_size512
intermediate_size2048
num_attention_heads8
vocab_size50304
max_position_embeddings2048
hidden_actgelu
layer_norm_eps1e-05
tie_word_embeddingsFalse
dtypefloat16
use_parallel_residualTrue
attention_biasTrue

The Notes

03 · what to know before you run

The block is parallel on Pythia

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

Both LayerNorms read the block input x, each with its own weight and bias, so self_attn.input and mlp.input are two different normalizations of the same tensor. The name post_attention_layernorm is the trap: on this block it does not follow the attention. Every Pythia config sets use_parallel_residual: true, and EleutherAI/gpt-neox-20b leaves it out and takes the config default, also true.

With use_parallel_residual: false the same modules run in sequence, as on Llama:

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

togethercomputer/RedPajama-INCITE-Base-3B-v1 is one such checkpoint. The diagram draws Pythia's setting; read model.config.use_parallel_residual before reusing a recipe on another GPTNeoXForCausalLM checkpoint.

A block's MLP never sees its own attention

The contribution identity holds in both settings: attention_output is the attention module's output (dense, bias included) and mlp_output the MLP's. On the parallel block nothing connects the two sublayers of one block: an edit to attention_output leaves that block's mlp.input bit-identical, and the first MLP that can read block i's attention is block i + 1's. Circuit and path-patching graphs on Pythia have no attention-to-MLP edge inside a block.

with model.trace(prompt):
    x = model.layers[2].input.save()
    attn = model.layers[2].self_attn.attention_output.save()
    mlp_in = model.layers[2].mlp.input.save()
    mlp = model.layers[2].mlp.mlp_output.save()
    out = model.layers[2].layer_output.save()

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

with model.trace(prompt):
    model.layers[2].self_attn.attention_output[:] = 0
    mlp_in_ablated = model.layers[2].mlp.input.save()

assert torch.equal(mlp_in, mlp_in_ablated)

The block returns a bare tensor, so model.layers[i].output is layer_output. post_attention_layernorm still runs after the attention returns, so in one trace read the attention's values before mlp.input.

Queries, keys and values come from one projection, laid out per head

self_attn.query_key_value is a single Linear from hidden_size to 3 * hidden_size, with a bias. Its output is grouped by head, [q_h | k_h | v_h] for each head h in turn, not [Q | K | V] in three blocks: the first hidden_size columns are not the queries. Split it by viewing it as [batch, seq, num_heads, 3 * head_dim]. Keys and values have num_heads heads; there is no grouping.

attention_queries and attention_keys are read after the rotary embedding, attention_values as split. For the queries and keys before the rotary, read the projection's output in the same trace:

with model.trace(prompt):
    qkv = model.layers[2].self_attn.source.self_query_key_value_0.output.save()
    q = model.layers[2].self_attn.attention_queries.save()

b, s, _ = qkv.shape
per_head = qkv.view(b, s, model.num_heads, 3 * model.head_dim).transpose(1, 2)
q_pre, k_pre, v = per_head.chunk(3, dim=-1)
rot = model.head_dim // 4                            # partial_rotary_factor 0.25
assert torch.equal(q_pre[..., rot:], q[..., rot:])   # past rot: no rotation

Rotary turns a quarter of each head

The Pythia configs set rotary_pct: 0.25, which transformers reads as rope_parameters["partial_rotary_factor"]. On 70m-deduped that is the first 16 of each head's 64 dimensions, rotated as two halves of 8 (rotate_half); the other 48 dimensions of every query and key are the projection's output unchanged. Only those 16 dimensions make a query-key score depend on the distance between the two tokens; the other 48 add the same term at any distance. The scale is head_dim ** -0.5, and attention_scores are the scaled, masked scores at the softmax's input, so softmax(attention_scores) equals attention_probabilities. Other GPTNeoXForCausalLM checkpoints set other fractions (RedPajama-INCITE-Base-3B-v1 rotates whole heads, 1.0).

Load with eager for the interior

Under the default sdpa load the six interior values report unavailable, and under attn_implementation="eager" all are served. The two implementations compute the same attention (no softcap, sink or window here), so the default load is the same model up to rounding. The Pythia files are float16, and a load without dtype keeps that, on CPU too; pass dtype=torch.float32 for exact comparisons.

model = StandardizedTransformer(
    "EleutherAI/pythia-70m-deduped", dispatch=True, attn_implementation="eager"
)

The final LayerNorm's bias is a fixed logit vector

logits equals lm_head.output, and project_on_vocab applied to the last block's layer_output equals logits exactly. The final norm is a LayerNorm with a bias, so a logit lens adds lm_head.weight @ norm.bias at every block, whatever the input. On 70m-deduped that vector is large, a standard deviation of 41 logits against 47 for a whole lens readout at block 3 on a sample prompt, and its top tokens are ,, ., \n, that, to. Folding the norm into the unembedding has to keep it as a bias term.

bias_logits = model.lm_head.weight @ model.norm.bias
top = [model.tokenizer.decode(i) for i in bias_logits.topk(5).indices.tolist()]

embed_in and lm_head are separate matrices, so the unembedding is not the embedding's transpose and an edit to one leaves the other. The vocabulary is padded: 50304 rows on 70m-deduped against 50277 tokenizer tokens. The 27 extra ids decode to ''; their unembedding rows are not zero, so they get logits.

The tokenizer adds no BOS

A prompt is tokenized without a leading token: position 0 is the prompt's first token, not a fixed BOS. <|endoftext|> (id 0) is both bos_token and eos_token, and <|padding|> (id 1) is the pad token.

The Pythia suite: sizes, deduplication, training steps

Pythia is ten sizes, 14m to 12b, each trained on the Pile and, as -deduped, on the Pile after global deduplication; the model cards state that every model saw the same data in the same order. Every config sets the parallel block, rotary_pct: 0.25 and untied embeddings. Each repo has 154 training checkpoints as branches, step0, step1, step2, step4 ... step512, then every 1000 steps to step143000; load one with revision. EleutherAI also publishes retrained seeds (pythia-70m-seed1, ...) and the -v0 runs.

model = StandardizedTransformer(
    "EleutherAI/pythia-70m-deduped", revision="step1000", dispatch=True
)

Other checkpoints load as GPTNeoXForCausalLM with different settings: EleutherAI/gpt-neox-20b (parallel, gelu_fast), stabilityai/stablelm-base-alpha-3b (parallel), togethercomputer/RedPajama-INCITE-Base-3B-v1 (sequential, full rotary).

Sparse autoencoders and tuned lenses

  • saprmarks/pythia-70m-deduped-saes, the dictionaries of Sparse Feature Circuits: embed, attn_out_layerN, mlp_out_layerN and resid_out_layerN read the outputs of embed_in, the attention, the MLP and the block, which are model.token_embeddings, model.layers[N].self_attn.attention_output, model.layers[N].mlp.mlp_output and model.layers[N].layer_output.
  • EleutherAI/sae-pythia-70m-deduped-32k and EleutherAI/sae-pythia-160m-deduped-32k: the folders layers.N.attention, layers.N.mlp and layers.N are SAEs on those modules' outputs, attention_output, mlp_output and layer_output.
  • Tuned lenses for the deduped Pythia models and GPT-NeoX-20B, in the AlignmentResearch/tuned-lens Hugging Face space that the tuned-lens package loads from. A lens for block N reads transformers' hidden_states[N], which is model.layers[N].input, and decodes through the model's own final norm and unembedding, as project_on_vocab does.

On 70m-deduped at block 3, both sets of SAEs reconstruct the value named above (fraction of variance unexplained 0.02 to 0.39) and not the sublayer's input (self_attn.input, mlp.input: above 1).

From the family module

04 · nnterp/families/gpt_neox.py

GPT-NeoX (GPTNeoXForCausalLM; Pythia, GPT-NeoX-20B).

Its tree is gpt_neox.{embed_in, emb_dropout, layers[i].{input_layernorm, post_attention_layernorm, attention, mlp}, final_layer_norm}. The head is lm_head on current transformers and was embed_out before; that key binds only where it resolves, so both spellings map to lm_head.