nnterp · encyclopedia granite · nnterp.families.granite

family granite

Granite

Llama's block with four config scalars around it: the embeddings are multiplied before block 0, each sublayer's output is multiplied by residual_multiplier before it is added, the attention scales its scores by attention_multiplier, and the logits are the head's output divided by logits_scaling.

Scaled residual addsEmbedding multiplierScaled logitsGrouped-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× 40
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsscaled
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 GraniteForCausalLMlayoutwhat 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] GraniteDecoderLayerlayoutwhat 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 GraniteAttentionlayoutwhat 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: its output times residual_multiplier

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 GraniteMLPlayoutwhat 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: its output times residual_multiplier

the module's own output

Printout

What print(model) shows on ibm-granite/granite-3.0-2b-base: the native modules, standard names and every standard value with its layout.

GraniteForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=49152, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(49152, 2048, padding_idx=0)
  (layers): ModuleList(
    (0-39): 40 x GraniteDecoderLayer(
      (self_attn): GraniteAttention(
        (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (k_proj): Linear(in_features=2048, out_features=512, bias=False)
        (v_proj): Linear(in_features=2048, out_features=512, bias=False)
        (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (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_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): GraniteMLP(
        (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: its output times residual_multiplier
        (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
      )
      (input_layernorm): GraniteRMSNorm((2048,), eps=1e-05)
      (post_attention_layernorm): GraniteRMSNorm((2048,), eps=1e-05)
      (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): GraniteRMSNorm((2048,), eps=1e-05)
  (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_layers40
model.hidden_size2048
model.num_heads32
model.num_kv_heads8
model.head_dim64
model.qk_head_dim64
model.vocab_size49152
model.intermediate_size8192
architecturesGraniteForCausalLM
num_hidden_layers40
hidden_size2048
intermediate_size8192
num_attention_heads32
num_key_value_heads8
vocab_size49152
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
embedding_multiplier12.0
residual_multiplier0.22
attention_multiplier0.015625
logits_scaling8.0
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

x      = embed_tokens(ids) * embedding_multiplier
h      = x + self_attn(input_layernorm(x)) * residual_multiplier
out    = h + mlp(post_attention_layernorm(h)) * residual_multiplier
logits = lm_head(norm(out)) / logits_scaling

The modules and their names are Llama's, post_attention_layernorm included (the MLP's input norm). The four scalars sit between the modules and are read from the config, so each is a plain number, the same on every block. On the released checkpoints:

  • 3.x 2B: embedding_multiplier 12, residual_multiplier 0.22, attention_multiplier 1/64, logits_scaling 8.
  • 3.x 8B and 4.1-8B: 12, 0.22, 1/128, 16.
  • 4.1-3B: 12, 0.22, 1/64, 10.
  • 4.1-30B: 12, 0.175, 1/128, 16.
  • 4.2 (3B, 8B, 30B): 1, 1, 1 / head_dim, 1, with lm_head untied. On 4.2 the contributions equal the modules' outputs and logits equals lm_head.output.

The contributions are scaled copies of the modules' outputs

attention_output is self_attn.output[0] * residual_multiplier and mlp_output is mlp.output * residual_multiplier: the tensors the block adds, computed from the modules' outputs, which are not themselves on the stream. With them the plain identity holds bit for bit on 2B, in float32 and bfloat16:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn_raw = model.layers[1].self_attn.output[0].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()

m = model.config.residual_multiplier                  # 0.22
torch.testing.assert_close(attn, attn_raw * m)
torch.testing.assert_close(x + attn + mlp, out)

Direct logit attribution and any decomposition of the stream into sublayer terms use attention_output and mlp_output; self_attn.output[0] and mlp.output are 1/0.22 ≈ 4.5 times larger than what reaches the stream. A vector added to mlp_output reaches the stream as it is; the same vector added to mlp.output reaches it times 0.22. A read changes nothing: the forward stays bit-identical.

A write to a contribution rounds every position

A write to attention_output or mlp_output reaches the model by dividing the whole edited copy by residual_multiplier and handing it back as the module's output, and the block multiplies it again. Positions you did not edit go through that round trip too. On 2B in bfloat16, adding a unit-norm vector to mlp_output at the last position of block 1 moved the logits at every earlier position by up to 0.19 (a KL of up to 0.002), although the edit cannot reach those positions causally. The edit's own effect at the edited position was a KL of 8e-5. In float32 the stray change is 2e-5 in the logits. Editing the module's output by the vector over the multiplier gives the same effect at the edited position, up to rounding, and leaves every other position bit-identical in either dtype:

m = model.config.residual_multiplier
with model.trace(prompt):
    model.layers[1].mlp.mlp_output[:, -1] += v    # every position rounds

with model.trace(prompt):
    model.layers[1].mlp.output[:, -1] += v / m    # position -1 only

For fine-grained edits (one position, one head, a small steering vector), edit the module's output this way or load in float32. The 3.0 base checkpoints (2B, 8B) are stored in float32, the others in bfloat16.

The query scale is attention_multiplier, 1/head_dim

The attention multiplies q · kᵀ by config.attention_multiplier in place of head_dim ** -0.5. On every released checkpoint it is 1 / head_dim (1/64 on 2B), an eighth of Llama's scale at head_dim 64. attention_scores are read after the scale and the causal mask, and softmax(attention_scores) reproduces attention_probabilities exactly in float32. Code that recomputes scores from queries and keys takes the multiplier from the config:

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

k = k.repeat_interleave(model.num_heads // model.num_kv_heads, dim=1)
scale = model.config.attention_multiplier    # not head_dim ** -0.5
qk = q @ k.transpose(-1, -2) * scale
causal = torch.ones_like(qk[0, 0], dtype=torch.bool).tril()
torch.testing.assert_close(qk[..., causal], scores[..., causal])

The 2B has 32 query heads over 8 key/value heads: attention_keys and attention_values are [batch, 8, seq, 64], read before repeat_kv, and an edit to key/value head j reaches query heads 4j to 4j + 3. A default load runs sdpa, which computes the same function: on 2B its logits equal eager's bit for bit in float32, and differ by up to 0.23 in bfloat16 with the same top token. The six interior values need attn_implementation="eager".

Logits are the head's output divided by logits_scaling

model.logits is model.lm_head.output / logits_scaling, exactly, and project_on_vocab divides too, so a logit lens at the last block equals logits. A softmax over lm_head.output, or over lm_head(norm(hidden)) written by hand, is at temperature 1/8 on 2B: on a test prompt the top token's probability is 0.94 from logits and 1.00 from lm_head.output. Take probabilities and KLs from logits or project_on_vocab.

with model.trace(prompt):
    resid = model.layers[-1].layer_output.save()
    raw = model.lm_head.output.save()
    logits = model.logits.save()

torch.testing.assert_close(logits, raw / model.config.logits_scaling)
torch.testing.assert_close(model.project_on_vocab(resid), logits)

lm_head and embed_tokens are one parameter on 3.x and 4.1, so an edit to embed_tokens.weight is an edit to the unembedding.

token_embeddings is the unscaled lookup

The embedding module returns the plain lookup, and the model multiplies it by embedding_multiplier afterwards, so token_embeddings is twelve times smaller than layers[0].input (row norms about 0.9 against 10.7 on 2B). Embeddings passed as inputs_embeds are multiplied too. An edit to token_embeddings reaches block 0 times 12; to set what block 0 reads, edit layers[0].input.

with model.trace(prompt):
    emb = model.token_embeddings.save()
    first = model.layers[0].input.save()

torch.testing.assert_close(first, emb * model.config.embedding_multiplier)

Position 0 is a sink, and no BOS token is added

Neither tokenizer prepends a beginning-of-sequence token, so the first real token takes the role. On 2B its residual norm is 293 after block 2, against 8 to 12 at the other positions of a test prompt, and about 4000 after block 20, against a mean of 16; in blocks 5, 10 and 20 heads put 0.71 to 0.83 of their attention on it (two test prompts). 4.1-3B shows the same at block 5 (norm 107, 0.67 of the attention). A later token can grow too: the reached 243 after block 5 in one prompt. Slice [:, 1:] before averaging activations for steering vectors, mean ablation or probes.

with model.trace(prompt):
    resid = model.layers[2].layer_output.save()

resid[0].norm(dim=-1)        # 293 at position 0, 8 to 12 elsewhere on 2B

Target tokens differ by tokenizer

The 3.x tokenizer (49152 tokens) splits " Paris" into ĠPar and is; the 4.1 tokenizer (100352 tokens) has " Paris" (12366) and "Paris" (60704) as whole tokens. Decode the target ids on the checkpoint you run.

What this family module covers

The module serves every checkpoint whose config says model_type granite: the dense Granite 3.0, 3.1, 3.2 and 3.3 at 2B and 8B, and Granite 4.1 and 4.2 at 3B, 8B and 30B. Granite-7B and the Granite Code 3B and 8B are llama, Granite Code 20B and 34B gpt_bigcode. The mixtures of experts (3.x 1B-A400M, 3B-A800M) are granitemoe. All of Granite 4.0 is granitemoehybrid, including the attention-only Micro, 1B and 350M. The Swash models are granite_swa and granitemoe_swa.

From the family module

04 · nnterp/families/granite.py

Granite (GraniteForCausalLM, IBM Granite 3.x / 4.x dense).

Llama's tree and names, with four scalars from the config applied around it:

  • embedding_multiplier: the model multiplies the embedding module's output by it before the first block, so token_embeddings (the module's output) times the multiplier is layers[0].input.
  • attention_multiplier: the attention's softmax scale in place of 1/sqrt(head_dim), passed to the shared interface; nothing to override.
  • residual_multiplier: the block adds h * residual_multiplier for each sublayer, not h. What reaches the residual stream is the scaled term, so attention_output and mlp_output are the module's output times the multiplier. That product is a computed copy, so each value divides on assignment, and a transform divides an edited copy back into the module's output so in-place edits reach the model.
  • logits_scaling: the model divides the head's output by it, so the family defines project_on_vocab with that step and the logit lens on the last block equals logits.