nnterp · encyclopedia starcoder2 · nnterp.families.starcoder2

family starcoder2

StarCoder2

Llama's block with LayerNorms, a bias on every projection, a two-matrix GELU MLP (c_fc, c_proj) in place of SwiGLU, and a sliding window of 4096 positions on every block.

LayerNormBiased q, k, vSliding windowGrouped-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× 30
→
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 Starcoder2ForCausalLMlayoutwhat 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] Starcoder2DecoderLayerlayoutwhat 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 Starcoder2Attentionlayoutwhat 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 Starcoder2MLPlayoutwhat 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/starcoder2-3b: the native modules, standard names and every standard value with its layout.

Starcoder2ForCausalLM(
  (lm_head): Linear(in_features=3072, out_features=49152, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(49152, 3072)
  (layers): ModuleList(
    (0-29): 30 x Starcoder2DecoderLayer(
      (self_attn): Starcoder2Attention(
        (q_proj): Linear(in_features=3072, out_features=3072, bias=True)
        (k_proj): Linear(in_features=3072, out_features=256, bias=True)
        (v_proj): Linear(in_features=3072, out_features=256, bias=True)
        (o_proj): Linear(in_features=3072, out_features=3072, 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): Starcoder2MLP(
        (c_fc): Linear(in_features=3072, out_features=12288, bias=True)
        (down_proj/c_proj): Linear(in_features=12288, out_features=3072, bias=True)
        (act_fn/act): GELUTanh()
        (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
      )
      (input_layernorm): LayerNorm((3072,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((3072,), eps=1e-05, elementwise_affine=True, bias=True)
      (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((3072,), 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_layers30
model.hidden_size3072
model.num_heads24
model.num_kv_heads2
model.head_dim128
model.qk_head_dim128
model.vocab_size49152
model.intermediate_size12288
architecturesStarcoder2ForCausalLM
num_hidden_layers30
hidden_size3072
intermediate_size12288
num_attention_heads24
num_key_value_heads2
vocab_size49152
max_position_embeddings16384
sliding_window4096
hidden_actgelu_pytorch_tanh
tie_word_embeddingsTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))     # LayerNorm, biased q, k, v, o
out = h + mlp(post_attention_layernorm(h))  # LayerNorm, c_fc -> GELU -> c_proj

Llama's tree and Llama's names, with LayerNorms (a mean subtracted, a weight and a bias) where Llama has RMSNorms. Every linear map carries a bias (use_bias): q_proj, k_proj, v_proj, o_proj, c_fc and c_proj. The MLP has no gate: c_fc widens the normed stream four times, GELU (the tanh approximation, gelu_pytorch_tanh) acts on it, and c_proj maps it back, so a neuron is one column of c_fc and one row of c_proj. The attention and the MLP each end in a residual dropout inside the module, the identity in eval.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output, added to the stream with nothing in between. The identity holds exactly, in float32 on the pinned checkpoint and in bfloat16 on starcoder2-3b (block 5):

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

torch.equal(x + attn + mlp, out)   # True

With the biases, a zeroed sublayer input does not give a zero contribution: the attention falls back to its value and output biases, and the MLP returns c_proj(gelu(c_fc.bias)), which includes c_proj.bias. To remove a sublayer's effect, ablate attention_output or mlp_output itself.

Every block attends over a window of 4096

config.sliding_window is 4096 on every StarCoder2 checkpoint, against a max_position_embeddings of 16384, and the config has no layer_types: every block takes the window. A query attends to its own position and the 4095 before it, so on a prompt shorter than 4096 tokens the mask is the plain causal mask. The window is in the mask the attention receives, so under attn_implementation="eager" it shows in attention_scores (the largest negative float outside it) and in attention_probabilities (zero outside it). On the pinned checkpoint loaded with sliding_window=4, each query has at most 4 nonzero keys, itself and the 3 before:

model = StandardizedTransformer(repo, attn_implementation="eager", sliding_window=4)
with model.trace(prompt):
    p = model.layers[1].self_attn.attention_probabilities.save()

(p[0, 0] > 0).sum(-1)   # 1, 2, 3, 4, 4, 4, ...

The pinned checkpoint itself has sliding_window null.

Grouped-query attention, few key/value heads

Every size has head_dim 128. 3B has 24 query heads over 2 key/value heads, 7B 36 over 4, 15B 48 over 4. On 3B attention_keys and attention_values are [batch, 2, seq, 128] and query head h reads key/value head h // 12: an edit to key/value head j reaches twelve query heads, 12j to 12j + 11, half of the block's heads. The query scale is head_dim ** -0.5; the rotary is the plain one over the whole head, with rope_theta near 1000000 on 3B and 7B and 100000 on 15B.

Load with eager for the interior

A default load runs sdpa, and the six attention interior values report unavailable; attn_implementation="eager" serves them. Both paths apply the same window, so they compute the same function.

The readout and the tokenizer

model.logits is model.lm_head.output, and project_on_vocab is lm_head(norm(hidden)), the final LayerNorm's bias included. That bias adds the fixed vector lm_head.weight @ norm.bias to every lens readout. On 3B the vector is small: its standard deviation is 0.23 logits, against 2.15 for a whole readout at block 14 on a code prompt. On 3B and 7B lm_head is tied to embed_tokens (tie_word_embeddings defaults to true and the files store no lm_head weight), so an edit to embed_tokens.weight is an edit to the unembedding; 15B sets it to false.

The tokenizer has 49152 ids and prepends nothing: position 0 is the text's first token. The vocabulary is built for code, and common English words split: " Paris" is two tokens, " Par" and "is". <fim_prefix>, <fim_suffix> and <fim_middle> are special tokens for fill-in-the-middle prompts.

From the family module

04 · nnterp/families/starcoder2.py

StarCoder2 (Starcoder2ForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward. The norms are nn.LayerNorm; the MLP is c_fc/c_proj with no gate; the attention and MLP outputs pass a residual dropout inside their module, the identity in eval.