GPT-J

GPTJForCausalLM( (lm_head): Linear(in_features=4096, out_features=50400, bias=True) (generator): Generator( (streamer): Streamer() ) (embed_tokens): Embedding(50400, 4096) (layers): ModuleList( (0-27): 28 x GPTJBlock( (input_layernorm/ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True) (self_attn/attn): GPTJAttention( (attn_dropout): Dropout(p=0.0, inplace=False) (resid_dropout): Dropout(p=0.0, inplace=False) (k_proj): Linear(in_features=4096, out_features=4096, bias=False) (v_proj): Linear(in_features=4096, out_features=4096, bias=False) (q_proj): Linear(in_features=4096, out_features=4096, 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 ) (mlp): GPTJMLP( (fc_in): Linear(in_features=4096, out_features=16384, bias=True) (down_proj/fc_out): Linear(in_features=16384, out_features=4096, 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 ) (layer_mid): Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP: the one norm feeds both sublayers (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_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=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 )

One LayerNorm feeds both sublayers

h   = ln_1(x)
out = self_attn(h) + mlp(h) + x

ln_1 is input_layernorm and the block's only norm: there is no post_attention_layernorm. Its one output tensor is passed to the attention and then to the MLP, so self_attn.input and mlp.input are one tensor object, not two equal copies. The block returns (hidden_states, attn_weights) on every call: model.layers[i].output[1] is the attention pattern, equal to self_attn.attention_probabilities, and layer_output is the first element.

An in-place edit of `self_attn.input` reaches the MLP

Because both sublayers read one tensor, an in-place edit of self_attn.input also changes what the MLP reads. An assignment replaces only the attention's argument and leaves mlp.input as it was. An edit of mlp.input, in place or assigned, reaches only the MLP: the attention has already run. To change what both read, assign input_layernorm.output.

with model.trace(prompt):
    model.layers[4].self_attn.input[:, -1] = 0      # in place
    mlp_in = model.layers[4].mlp.input.save()       # its last row is zero too

with model.trace(prompt):
    # an assignment replaces only the attention's argument
    model.layers[4].self_attn.input = torch.zeros_like(model.layers[4].self_attn.input)
    mlp_in = model.layers[4].mlp.input.save()       # unchanged

The attention runs before the MLP, so inside one trace read and edit the attention's values before mlp.input; asking for mlp.input first makes a later attention read fail as out of order.

The contributions, added in the block's order

attention_output is the attention module's output (out_proj, no bias) and mlp_output the MLP's (fc_out, with a bias). The block computes attn + mlp + x, so that sum reproduces layer_output bit for bit; x + attn + mlp differs from it by rounding (up to 0.5 at block 4 of 6B in bfloat16 on a sample prompt). Nothing inside a block connects the two sublayers: zeroing attention_output leaves that block's mlp.input bit-identical, and the first MLP that reads block i's attention is block i + 1's.

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

assert torch.equal(attn + mlp + x, out)                 # the block's own order

The attention is eager by default and computes its scores in float32

transformers has no sdpa path for GPT-J: a load without attn_implementation runs the eager GPTJAttention, and all six interior values are served with no flag. They are read inside the eager attention, so a flash_attention_2 load reports them unavailable. The values are read around the module's _attn(query, key, value, mask) call: its arguments are attention_queries, attention_keys and attention_values (16 heads each on 6B, no grouping, in the model's dtype), and its first return is attention_head_outputs.

_attn casts the queries and keys to float32, so attention_scores is float32 whatever the model's dtype: q @ kᵀ / 16 (sqrt(head_dim), 256 on 6B) plus the causal mask, whose masked entries hold the model dtype's minimum (-3.39e38 under bfloat16). The pattern is cast back to the model's dtype, and softmax(attention_scores) cast to it equals attention_probabilities exactly. q_proj, k_proj, v_proj and out_proj are separate Linears with no bias; attention_head_outputs.flatten(-2) is out_proj.input.

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

again = q.float() @ k.float().transpose(-1, -2) / model.head_dim ** 0.5
causal = torch.ones(q.shape[2], q.shape[2], dtype=torch.bool, device=q.device).tril()
assert torch.equal(again[..., causal], scores[..., causal])

Rotary turns the first 64 dimensions, in adjacent pairs

rotary_dim is 64 of each head's 256 dimensions on 6B. attention_queries and attention_keys are read after the rotary; their dimensions 64 to 255 are the projection's output unchanged, and only the first 64 make a score depend on the distance between two tokens. Within those 64, GPT-J rotates adjacent pairs: dimensions 2i and 2i + 1 turn together at frequency 10000 ** (-2i / 64) (rotate_every_two). Llama and GPT-NeoX pair dimension i with i + d / 2 across a rotated span of width d (rotate_half). Query dimension 1 on GPT-J is the partner of dimension 0, so a comparison of queries or keys across families, or rotary code written for rotate_half, needs the rotated span reordered, evens then odds. Reordered, GPT-J's rotation is rotate_half's with the same frequencies; recomputed in float32 the two agree exactly.

rot = model.config.rotary_dim                                        # 64
perm = torch.cat([torch.arange(0, rot, 2), torch.arange(1, rot, 2)])

with model.trace(prompt):
    q = model.layers[4].self_attn.attention_queries.save()

q_half = torch.cat([q[..., :rot][..., perm], q[..., rot:]], dim=-1)   # rotate_half's

The sine and cosine table is a float32 buffer, cast to the model's dtype before it multiplies, so the queries and keys are served in the model's dtype.

Loading: the main branch is float32

The main branch of EleutherAI/gpt-j-6b holds float32 weights (24 GB) and its config names no dtype, so a load without dtype is float32. Pass a half dtype to fit a single GPU, or load the float16 branch, a 12 GB half-precision copy:

model = StandardizedTransformer(
    "EleutherAI/gpt-j-6b", dtype=torch.bfloat16, device="cuda", dispatch=True
)

The head has a bias, and so does the final norm

logits equals lm_head.output, and project_on_vocab applied to the last block's layer_output equals logits exactly. ln_f is a LayerNorm with a bias and lm_head adds a bias of its own, so a logit lens adds lm_head.weight @ norm.bias + lm_head.bias at every block, whatever the input. On 6B that vector has a standard deviation of 0.40 logits, most of it from the norm's bias (lm_head.bias alone: 0.034), and its top tokens are ,, \n, ., -, and. Folding the norm into the unembedding, or a direct logit attribution, keeps both as constant terms. wte and lm_head are separate matrices: the unembedding is not the embedding's transpose.

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

The vocabulary is padded to 50400

The tokenizer is GPT-2's BPE, 50257 tokens, plus 143 added <|extratoken_N|> tokens (ids 50257 to 50399), so len(model.tokenizer) equals vocab_size, 50400. The model card states that only the 50257 GPT-2 ids are used. Their unembedding rows are smaller (norm 0.32 against 1.31) but not zero: on six sample prompts their logits stay below −3.7 while every position's top logit is above 12. The tokenizer adds no BOS: position 0 is the prompt's first token, and <|endoftext|> (50256) is the BOS, EOS and unknown token.

Model editing and causal tracing

ROME's GPT-J hyperparameters rewrite transformer.h.5.mlp.fc_out, which is model.layers[5].mlp.fc_out.weight. Its key is the input of that Linear at the subject's last token, model.layers[5].mlp.fc_out.input, and its value the output, which equals model.layers[5].mlp.mlp_output. MEMIT spreads the same update over blocks 3 to 8.
ROME's causal tracing on EleutherAI/gpt-j-6B adds noise to the output of transformer.wte, model.token_embeddings, and restores the outputs of transformer.h.{i}, .mlp and .attn (first element), which are layer_output, mlp_output and attention_output.
Function vectors (Todd et al.) on GPT-J read each head from the input of attn.out_proj split by head, which is model.layers[i].self_attn.attention_head_outputs, and add the vector to the output of transformer.h.{i}, model.layers[i].layer_output.

The block

The API

Values, by host

Printout

Config

Config

Config

The Notes

One LayerNorm feeds both sublayers

An in-place edit of `self_attn.input` reaches the MLP

The contributions, added in the block's order

The attention is eager by default and computes its scores in float32

Rotary turns the first 64 dimensions, in adjacent pairs

Loading: the main branch is float32

The head has a bias, and so does the final norm

The vocabulary is padded to 50400

Model editing and causal tracing

From the family module

GPT-J

The block

The API

Values, by host

Printout

Config

Config

Config

The Notes

One LayerNorm feeds both sublayers

An in-place edit of self_attn.input reaches the MLP

The contributions, added in the block's order

The attention is eager by default and computes its scores in float32

Rotary turns the first 64 dimensions, in adjacent pairs

Loading: the main branch is float32

The head has a bias, and so does the final norm

The vocabulary is padded to 50400

Model editing and causal tracing

From the family module

An in-place edit of `self_attn.input` reaches the MLP