nnterp · encyclopedia stablelm · nnterp.families.stablelm

family stablelm

StableLM 2StableLM-3B

Llama's block with LayerNorms and rotary on the first quarter of each head; StableLM-2-12B makes it parallel (use_parallel_residual) and norms each query and key head (qk_layernorm).

LayerNormPartial rotaryBiased q, k, v

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 StableLmForCausalLMlayoutwhat 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] StableLmDecoderLayerlayoutwhat 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 StableLmAttentionlayoutwhat 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 StableLmMLPlayoutwhat 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 stabilityai/stablelm-2-1_6b: the native modules, standard names and every standard value with its layout.

StableLmForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 2048)
  (layers): ModuleList(
    (0-23): 24 x StableLmDecoderLayer(
      (self_attn): StableLmAttention(
        (q_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (k_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (v_proj): Linear(in_features=2048, out_features=2048, bias=True)
        (o_proj): Linear(in_features=2048, out_features=2048, 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_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): StableLmMLP(
        (gate_proj): Linear(in_features=2048, out_features=5632, bias=False)
        (up_proj): Linear(in_features=2048, out_features=5632, bias=False)
        (down_proj): Linear(in_features=5632, out_features=2048, bias=False)
        (act_fn): SiLUActivation()
        (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((2048,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True, bias=True)
      (dropout): Dropout(p=0.0, inplace=False)
      (layer_mid): The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (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((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_heads32
model.num_kv_heads32
model.head_dim64
model.qk_head_dim64
model.vocab_size100352
model.intermediate_size5632
architecturesStableLmForCausalLM
num_hidden_layers24
hidden_size2048
intermediate_size5632
num_attention_heads32
num_key_value_heads32
vocab_size100352
max_position_embeddings4096
hidden_actsilu
layer_norm_eps1e-05
tie_word_embeddingsFalse
dtypefloat16
use_parallel_residualFalse
qk_layernormFalse

The Notes

03 · what to know before you run

The block is sequential on 1.6B and 3B, parallel on 12B

# StableLM-2-1.6B, StableLM-3B-4E1T (use_parallel_residual false)
h   = x + self_attn(input_layernorm(x))
out = h + mlp(post_attention_layernorm(h))

# StableLM-2-12B (use_parallel_residual true)
n   = input_layernorm(x)
out = x + self_attn(n) + mlp(n)

The configs differ: only stablelm-2-12b sets use_parallel_residual: true, and the others leave it out and take the default, false. On the parallel block post_attention_layernorm does not exist, both sublayers read the one input_layernorm output, and an edit to attention_output leaves that block's mlp.input bit-identical. The diagram draws each checkpoint's own block: the parallel one on 12B, the sequential one on 1.6B and 3B. Read model.config.use_parallel_residual before reusing a recipe across sizes. The pinned checkpoint has 12B's shape.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output; the block's dropout after the MLP is the identity in eval. The identity holds exactly in float32 on both shapes: on the pinned parallel checkpoint and on StableLM-2-1.6B (block 5). On the sequential block the stream between the adds, post_attention_layernorm.input, equals layers[i].input + attention_output.

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

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

Rotary turns a quarter of each head

Every config sets partial_rotary_factor: 0.25. On 1.6B, head_dim is 64, so the first 16 dimensions of each query and key head are rotated, as two halves of 8 (rotate_half), and the other 48 are the projection's output unchanged; only those 16 make a score depend on the distance between two tokens. On StableLM-2-1.6B the queries past dimension 16 equal q_proj's output split into heads, and the first 16 do not:

attn = model.layers[5].self_attn
with model.trace(prompt):
    q_raw = attn.source.self_q_proj_0.output.save()
    q = attn.attention_queries.save()

b, s, _ = q_raw.shape
q_pre = q_raw.view(b, s, -1, model.head_dim).transpose(1, 2)
rot = int(model.head_dim * model.config.rope_parameters["partial_rotary_factor"])   # 16
torch.equal(q_pre[..., rot:], q[..., rot:])   # True

The query scale is head_dim ** -0.5, and rope_theta is 10000.

Per-head LayerNorms on queries and keys on 12B

stablelm-2-12b sets qk_layernorm: true: the attention has q_layernorm and k_layernorm, each a list of LayerNorms without bias, one per query head (32) and one per key/value head (8), applied to the split heads before the rotary. attention_queries and attention_keys are read after both, so past the rotary dimensions they equal q_layernorm.output and k_layernorm.output, not the projections (checked on the pinned checkpoint, which has them). The other checkpoints have no such modules.

Biased queries, keys and values on StableLM 2 1.6B

use_qkv_bias is true on stablelm-2-1_6b and its zephyr, false on 12B and the 3B models. Where set, q_proj, k_proj and v_proj add a bias before the rotary; o_proj and the MLP have none. On 1.6B the biases are large: the key bias has a larger norm than the input-dependent part W_k·x (median over the tokens of a prompt) on most blocks, 14.5 times on block 0, and the query bias is the larger on blocks 0 to 4. So a zeroed attention input does not silence the attention; ablate attention_output itself.

Grouped-query attention on 12B

1.6B and the 3B models have 32 query heads and 32 key/value heads, one per query head. 12B has 32 query heads over 8 key/value heads, so an edit to key/value head j reaches query heads 4j to 4j + 3.

Load with eager for the interior

A default load runs sdpa, and the six attention interior values report unavailable; attn_implementation="eager" serves them. There is no window, softcap or sink, so the two compute the same function. The 1.6B files are float16; pass dtype=torch.float32 for exact comparisons.

The readout and the tokenizer

model.logits is model.lm_head.output, and lm_head has its own weight on every checkpoint. The final norm is a LayerNorm with a bias, so a logit lens adds lm_head.weight @ norm.bias at every block; on 1.6B its standard deviation is 0.31 logits and its top tokens are the, and, a, to, that. StableLM 2 uses a 100289-token tokenizer (vocab_size 100352), StableLM-3B GPT-NeoX's (50277 tokens, vocab_size 50304); neither prepends a BOS, so position 0 is the text's first token.

From the family module

04 · nnterp/families/stablelm.py

StableLM / StableLM-2 (StableLmForCausalLM).

Llama's containers. With use_parallel_residual the block is parallel with one input_layernorm and a dropout; without it there is a post_attention_layernorm too. Under qk_layernorm the attention norms its queries and keys per head (q_layernorm / k_layernorm). Of the released checkpoints only StableLM-2-12B sets both (as does the pinned tiny); StableLM-2-1.6B and StableLM-3B-4E1T are sequential with no q/k norms.