nnterp · encyclopedia codegen · nnterp.families.codegen

family codegen

CodeGen

GPT-J's parallel block and adjacent-pair rotary, with the queries, values and keys cut from one projection laid out as four groups of [q | v | k].

Parallel blockPartial rotaryInterleaved rotaryLayerNorm

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× 20
→
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 CodeGenForCausalLMlayoutwhat 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] CodeGenBlocklayoutwhat 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 Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them

nowhere: the family does not serve it
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP: the one norm feeds both sublayers

the child module input_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 CodeGenAttentionlayoutwhat 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 then scaled, in fp32

inside the forward: the input of call_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
model.layers[i].mlp CodeGenMLPlayoutwhat 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 Salesforce/codegen-350M-mono: the native modules, standard names and every standard value with its layout.

CodeGenForCausalLM(
  (lm_head): Linear(in_features=1024, out_features=51200, bias=True)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(51200, 1024)
  (layers): ModuleList(
    (0-19): 20 x CodeGenBlock(
      (input_layernorm/ln_1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True, bias=True)
      (self_attn/attn): CodeGenAttention(
        (attn_dropout): Dropout(p=0.0, inplace=False)
        (resid_dropout): Dropout(p=0.0, inplace=False)
        (qkv_proj): Linear(in_features=1024, out_features=3072, bias=False)
        (o_proj/out_proj): Linear(in_features=1024, out_features=1024, 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 then scaled, 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
      )
      (mlp): CodeGenMLP(
        (fc_in): Linear(in_features=1024, out_features=4096, bias=True)
        (down_proj/fc_out): Linear(in_features=4096, out_features=1024, 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((1024,), 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_layers20
model.hidden_size1024
model.num_heads16
model.num_kv_heads16
model.head_dim64
model.qk_head_dim64
model.vocab_size51200
model.intermediate_size4096
architecturesCodeGenForCausalLM
vocab_size51200
tie_word_embeddingsFalse
dtypefloat16
n_layer20
n_embd1024
n_head16
n_positions2048
activation_functiongelu_new
layer_norm_epsilon1e-05
rotary_dim32

The Notes

03 · what to know before you run

The block is parallel, with one LayerNorm

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

ln_1 is the block's only norm. Its output is both self_attn.input and mlp.input, the same tensor; input_layernorm is its alias, and the block has no post_attention_layernorm. attention_output is out_proj's output and mlp_output is fc_out's, and nothing normalizes either before the add. Nothing connects the two sublayers of one block: zeroing attention_output leaves that block's mlp.input bit-identical, so block i's attention first reaches an MLP at block i + 1.

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()

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

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's attn_weights are returned on every forward, not only under output_attentions: model.layers[i].output[1] is the block's attention pattern, equal to self_attn.attention_probabilities.

qkv_proj holds four groups of queries, values and keys

self_attn.qkv_proj is one Linear from hidden_size to 3 * hidden_size without a bias (out_proj has none either; fc_in and fc_out have one). The forward cuts its output into mp_num = 4 groups, a constant in the source, and splits each group as [q | v | k]: values in the middle, keys last. Group c holds heads c * num_heads / 4 to (c + 1) * num_heads / 4 - 1, four heads per group on 350M. A split of the full width into three hidden_size blocks, in either order, gets only the first group's queries right. Keys and values have num_heads heads; there is no grouping of key/value heads.

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

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

b, s, _ = qkv.shape
H, D = model.num_heads, model.head_dim
groups = qkv.view(b, s, 4, 3, H // 4, D)            # 4 groups of [q | v | k]
q_pre, v_split, k_pre = (
    groups[:, :, :, j].reshape(b, s, H, D).transpose(1, 2) for j in range(3)
)
assert torch.equal(v_split, v)
rot = model.config.rotary_dim                       # 32 of 64 on 350M
assert torch.equal(q_pre[..., rot:], q[..., rot:])  # past rotary_dim: not rotated

The weight's rows are in the same order: qkv_proj.weight.view(4, 3, num_heads // 4, head_dim, hidden_size)[c, j, h] is head c * num_heads / 4 + h of the queries (j = 0), values (1) or keys (2). An in-place edit of attention_queries, attention_keys or attention_values reaches the scores or the head outputs: the three are the tensors _attn computes with.

Rotary turns adjacent pairs of the leading channels

rotary_dim is 32 of each head's 64 channels on 350M; the configs set 64 on every larger size, of 80 on 2B and of 256 on 6B and 16B. The rotation pairs adjacent channels (2i, 2i + 1) (GPT-J's rotate_every_two), not the two halves (i, i + rotary_dim / 2) that rotate_half pairs on Llama and GPT-NeoX; pair i turns by position · 10000^(-2i / rotary_dim). A de-rotation or a frequency analysis written for rotate_half pairs the wrong channels here. The channels past rotary_dim are the projection's output unchanged and carry no position.

The masked scores are the dtype's minimum divided by the scale

attn_implementation="sdpa" fails at load; _attn is the only path. The mask holds the minimum of the model's dtype, and it is added before the division by sqrt(head_dim), so a masked entry of attention_scores is about that minimum divided by sqrt(head_dim). On 350M (head_dim 64) that is about -8188 on the default float16 load and -4.25e37 on a float32 load: finite, and neither the float16 nor the float32 minimum a mask test may look for. softmax(attention_scores) cast to the model's dtype equals attention_probabilities exactly.

The checkpoints are float16; the queries are not

The configs set torch_dtype: float16 (2B-nl: float32), and a load without dtype keeps it, on CPU too. The rotary multiplies by a float32 sin/cos buffer, so on that load attention_queries is float32 while attention_keys is float16. The cache casts the keys back to the model's dtype, and a trace runs with the cache on; under use_cache=False the keys stay float32. attention_values and attention_probabilities are float16, attention_scores float32. q @ k.transpose(-1, -2) on the two raises; cast first:

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

recomputed = q @ k.float().transpose(-1, -2) / model.head_dim ** 0.5

recomputed equals scores wherever the mask lets a key through.

The load report lists each block's attn.causal_mask as unexpected: the files store a mask buffer that transformers' class does not have, and the mask is built per call.

The head has a bias, and a constant logit vector

logits equals lm_head.output, and project_on_vocab applied to the last block's layer_output equals logits exactly, ln_f's bias and lm_head's included. A logit lens therefore adds the same vector at every block, whatever the input:

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

On 350M-mono this vector has a standard deviation of 1.8 logits, against about 4 for a whole lens readout at block 5 on a sample prompt. Most of that spread comes from ln_f.bias (lm_head.bias alone: 0.11), and its top tokens are \n, ,, -, ., _. Folding the norm into the unembedding keeps both biases as one term. wte and lm_head are separate matrices. The vocabulary is padded: 51200 rows against 50295 tokenizer tokens on -mono (50257 on -nl). On 350M-mono the 905 extra rows are not zero (biases between -0.38 and -0.34); on a sample prompt their total probability at any position is below 4e-6.

On -multi and -mono, runs of whitespace are single tokens

The -multi and -mono tokenizers are GPT-2's 50257 tokens, at the same ids, plus 38: one token for each run of 2 to 31 spaces (ids 50286 down to 50257) and of 2 to 9 tabs (50294 down to 50287). The -nl tokenizers are GPT-2's alone, so the same code prompt has different tokens and positions on -nl than on the other two. On -multi and -mono a run of two or more spaces becomes one token and the word after it loses its leading space, so an indented statement starts with the bare word:

model.tokenizer.encode("    return x")   # [50284, 7783, 2124]: '    ', 'return', ' x'

GPT-2 encodes the same text as three ' ' and ' return' (1441). A logit-lens or patching target inside a block is the bare return (7783): after "def add(a, b):\n " 350M-mono puts 0.34 on it. After the newline that opens a function body the model predicts the indent itself: after "def add(a, b):\n" 350M-mono puts 0.80 on the four-space token.

The tokenizer adds no BOS: position 0 is the prompt's first token. <|endoftext|> (50256) is its bos_token, eos_token and unk_token; the config's bos_token_id is 1, which is the token ".

Three data variants at four sizes

Each size was released three times; the model cards and the paper (A Conversational Paradigm for Program Synthesis) describe the data:

  • -nl: trained on the Pile.
  • -multi: initialized from -nl and trained further on BigQuery GitHub code in C, C++, Go, Java, JavaScript and Python (119.2B tokens).
  • -mono: initialized from -multi and trained further on BigPython (71.7B tokens of Python).

The three share the architecture and the 51200-row embedding; the tokenizers differ (above). The sizes, from the configs:

  • 350M: 20 blocks, hidden 1024, 16 heads × 64, rotary_dim 32
  • 2B: 32 blocks, hidden 2560, 32 heads × 80, rotary_dim 64
  • 6B: 33 blocks, hidden 4096, 16 heads × 256, rotary_dim 64
  • 16B: 34 blocks, hidden 6144, 24 heads × 256, rotary_dim 64

CodeGen2 loads under this name, and its 1B and 3.7B run wrong

Salesforce/codegen2-1B, codegen2-3_7B, codegen2-7B and codegen2-16B set model_type: "codegen" and ship their own modeling code (auto_map). Without trust_remote_code, transformers builds its own CodeGenForCausalLM, nnterp wraps it, and support() reports every value available. But the 1B and 3.7B code cuts qkv_proj into 8 groups (mp_num = 8), not 4, so the class mixes columns of the queries, values and keys. On a 13-line Python sample codegen2-1B's loss is 8.93 as loaded and 0.34 after regrouping each qkv_proj.weight from 8 groups of [q | v | k] to 4. The 7B and 16B code uses 4 groups. With trust_remote_code=True the remote code fails to import: it imports transformers.onnx, which transformers 5.20 does not have. CodeGen2.5 (Salesforce/codegen25-7b-mono, ...) is a LlamaForCausalLM and loads as Llama.

From the family module

04 · nnterp/families/codegen.py

CodeGen (CodeGenForCausalLM).

transformer.{wte, drop, h[i].{ln_1, attn, mlp}, ln_f} and lm_head. GPT-J's parallel block: ln_1 feeds both sublayers and x + attn + mlp is summed in the block, which returns (hidden_states, attn_weights). The attention projects queries, values and keys (in that order) out of one fused qkv_proj laid out for four TPU cores, applies rotary embeddings to the first rotary_dim channels, and does its own arithmetic in _attn: an fp32 q @ k^T, the mask added, then the division by sqrt(head_dim), softmax, a cast to the values' dtype, attn_dropout, @ v. The model has no attention implementation but that one, so attn_implementation never routes around it and the interior needs no eager load. drop has no standard name.

CodeGen2 checkpoints (Salesforce/codegen2-*) declare model_type codegen, but their own code lays qkv_proj out in eight groups (mp_num = 8); loaded through transformers' class, which splits it in four, the queries, keys and values are cut wrong and nothing reports it.