nnterp · encyclopedia gpt_bigcode · nnterp.families.gpt_bigcode

family gpt_bigcode

StarCoderSantaCoder

GPT-2's sequential block with multi-query attention: one key head and one value head serve every query head, and the queries, keys and values are cut from one fused c_attn projection.

Position embeddingsLayerNormMulti-query attention

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× 24
→
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 GPTBigCodeForCausalLMlayoutwhat 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] GPTBigCodeBlocklayoutwhat 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].self_attn GPTBigCodeAttentionlayoutwhat 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 GPTBigCodeMLPlayoutwhat 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 bigcode/gpt_bigcode-santacoder: the native modules, standard names and every standard value with its layout.

GPTBigCodeForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=49280, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(49280, 2048)
  (layers): ModuleList(
    (0-23): 24 x GPTBigCodeBlock(
      (input_layernorm/ln_1): LayerNorm((2048,), eps=1e-05, elementwise_affine=True, bias=True)
      (self_attn/attn): GPTBigCodeAttention(
        (c_attn): Linear(in_features=2048, out_features=2304, bias=True)
        (o_proj/c_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (resid_dropout): Dropout(p=0.1, inplace=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_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
      )
      (post_attention_layernorm/ln_2): LayerNorm((2048,), eps=1e-05, elementwise_affine=True, bias=True)
      (mlp): GPTBigCodeMLP(
        (c_fc): Linear(in_features=2048, out_features=8192, bias=True)
        (down_proj/c_proj): Linear(in_features=8192, out_features=2048, bias=True)
        (act_fn/act): GELUTanh()
        (dropout): Dropout(p=0.1, 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_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((2048,), 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_layers24
model.hidden_size2048
model.num_heads16
model.num_kv_heads1
model.head_dim128
model.qk_head_dim128
model.vocab_size49280
model.intermediate_size8192
architecturesGPTBigCodeForCausalLM
num_key_value_heads1
vocab_size49280
tie_word_embeddingsTrue
n_layer24
n_embd2048
n_head16
n_inner8192
n_positions2048
activation_functiongelu_pytorch_tanh
layer_norm_epsilon1e-05

The Notes

03 · what to know before you run

The block, in order

h   = x + attn(ln_1(x))      # ln_1 = input_layernorm
out = h + mlp(ln_2(h))       # ln_2 = post_attention_layernorm

The block returns a bare tensor. attention_output is self_attn.output[0] (after c_proj) and mlp_output is mlp.output (after the MLP's c_proj); nothing sits between either and the stream, and layers[i].input + attention_output + mlp_output == layer_output is exact (bit-identical on the pinned checkpoint). Under an eager load the attention calls transformers' eager attention function, where nnterp reads the six interior values.

One key head and one value head serve every query head

Every released checkpoint sets multi_query: the attention projects num_heads query heads and a single key head and value head, which every query head reads. model.num_kv_heads is 1, attention_keys and attention_values are [batch, 1, seq, head_dim], and an edit to them reaches every head's scores or outputs at once. To change what one head reads, edit its queries, its scores or its pattern instead. Head h's scores are q[:, h] @ k[:, 0]ᵀ * head_dim ** -0.5 (0.0884 at head_dim 128); the softmax runs in float32 and is cast back to the queries' dtype.

with model.trace(prompt):
    q = model.layers[1].self_attn.attention_queries.save()   # num_heads heads
    k = model.layers[1].self_attn.attention_keys.save()      # one head
    scores = model.layers[1].self_attn.attention_scores.save()

n = scores.shape[-1]
causal = torch.ones(n, n, dtype=torch.bool, device=scores.device).tril()
again = q @ k.transpose(-1, -2) * model.head_dim ** -0.5   # k broadcasts
torch.testing.assert_close(again[..., causal], scores[..., causal])

The configs' attention_softmax_in_fp32 and scale_attention_softmax_in_fp32 are not read by transformers' forward: the eager softmax is float32 whatever they say.

Queries, keys and values are cut from one Linear

c_attn is an nn.Linear from hidden_size to hidden_size + 2 * head_dim, its weight stored [out, in]. Its output is laid out as all the query heads, then the one key head, then the one value head, and within the queries head h is columns h * head_dim to (h + 1) * head_dim:

E, d = model.hidden_size, model.head_dim
W_Q, W_K, W_V = model.layers[1].self_attn.c_attn.weight.split([E, d, d])
b_Q, b_K, b_V = model.layers[1].self_attn.c_attn.bias.split([E, d, d])

with model.trace(prompt):
    h = model.layers[1].self_attn.input.save()               # ln_1's output
    q = model.layers[1].self_attn.attention_queries.save()

q_again = (h @ W_Q.T + b_Q).unflatten(-1, (model.num_heads, d)).transpose(1, 2)
torch.testing.assert_close(q_again, q)

With multi_query false, transformers lays c_attn's output out per head instead, as [q | k | v] for each of num_heads heads.

Assign the queries; do not edit them in place

The queries are a view of the split c_attn output, so with gradients on torch refuses an in-place edit of attention_queries (RuntimeError: ... is a view and is being modified inplace). Edit a copy and assign it, or edit under torch.no_grad(). The keys and values take an in-place edit.

with model.trace(prompt):
    q = model.layers[1].self_attn.attention_queries.clone()
    q[:, 0] = 0                                              # head 0's queries
    model.layers[1].self_attn.attention_queries = q

Load with eager

The default load runs sdpa, and the six interior values are unavailable under it; attn_implementation="eager" serves them. attention_output and mlp_output are there under either load.

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 n_positions (2048 on SantaCoder, 8192 on StarCoder and tiny_starcoder_py), 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 tokenizer adds no BOS, so position 0 is the prompt's first token. On tiny_starcoder_py, on a 58-token Python function, blocks 12 to 19 put 54% to 80% of their pattern (averaged over heads and the later queries) on key 0.

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 tiny_starcoder_py). 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 (a standard deviation of 0.12 logits on tiny_starcoder_py).

The family's checkpoints

SantaCoder (gpt_bigcode-santacoder, 1.1B) is 24 blocks × 2048 with 16 query heads of 128 and gelu_pytorch_tanh. tiny_starcoder_py is 20 × 768 with 12 query heads of 64. StarCoder, StarCoderBase (with its 1B, 3B and 7B sizes) and StarCoderPlus are gated. Their fine-tunes OctoCoder, StarChat-alpha and WizardCoder-15B are open, and are 40 × 6144 with 48 query heads of 128 and gelu. Every one has an MLP four times the width (n_inner) and one key/value head. bigcode/santacoder is SantaCoder's original repository, a gpt2 config with remote code (GPT2LMHeadCustomModel), so it does not load as this family; gpt_bigcode-santacoder does.

From the family module

04 · nnterp/families/gpt_bigcode.py

GPT-BigCode (GPTBigCodeForCausalLM): StarCoder, SantaCoder.

GPT-2's tree and GPT-2's block: transformer.{wte, wpe, h[i].{ln_1, attn, ln_2, mlp}, ln_f} and lm_head, the residual added in the block, which returns a bare tensor. The attention projects queries, keys and values with one fused c_attn and runs the shared eager forward. Under multi_query (every released checkpoint) the keys and values are one head shared by every query head, so num_kv_heads is 1. Queries, keys and values are split views of the one c_attn tensor: torch refuses an in-place edit of attention_queries outside torch.no_grad() (assign instead), while keys and values take one. attention_softmax_in_fp32 and scale_attention_softmax_in_fp32 are stored on the module but not read: the eager softmax runs in float32 regardless. The MLP is n_inner wide, 4 * hidden_size when that is None.