GraniteForCausalLM

GraniteForCausalLM at 40 layers and hidden size 4096. 19 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From granite-4.1-8b

Also covers granite-4.1-8b-base, granite-guardian-4.1-8b, granite-3.3-8b-instruct, granite-3.3-8b-math-prm-v2, granite-guardian-3.3-8b, granite-3.1-8b-instruct, granite-3.1-8b-lora-intrinsics-v0.1, granite-3.2-8b-instruct, granite-3.2-8b-instruct-preview, granite-guardian-3.1-8b, granite-guardian-3.2-8b-factuality-detection, granite-3.1-8b-base, granite-3.3-8b-base.

x = embed[ids] * 12                    # [T, 4096]   <- embedding_multiplier

for i in 0 .. 39:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T       # [T, 32*128]
    k = h @ layers[i].self_attn.k_proj.T       # [T, 8*128]
    v = h @ layers[i].self_attn.v_proj.T       # [T, 8*128]
    q, k = rope(q, k, theta=10000000)
    k, v = repeat_kv(k, v, 4)            # 32 query heads share 8 key/value heads
    a = softmax(q @ k.T * 0.0078125, mask=causal) @ v      # <- attention_multiplier
    a = a @ layers[i].self_attn.o_proj.T
    x = x + a * 0.22        # <- residual_multiplier

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g = silu(h @ layers[i].mlp.gate_proj.T)           # [T, 12800]
    u =       h @ layers[i].mlp.up_proj.T
    y = (g * u) @ layers[i].mlp.down_proj.T
    x = x + y * 0.22

x = rmsnorm(x, model.norm)
logits = (x @ embed.T) / 16        # tied by the configuration (see Notes); <- logits_scaling

Differs from the block above

The same arithmetic with these values changed. Everything not listed is as above.

Notes

head_dim is not published; 4096 / 32 = 128 is used.

Output head: the configuration says tied, the checkpoint ships a separate one. granite-4.1-8b sets tie_word_embeddings: true, which is what the block follows, but its checkpoint also ships a separate lm_head.weight [100352, 4096]. A standard transformers load reads that tensor and then, because the configuration asks for tying, replaces it with the input embedding (modeling_utils.py L2994, L5065), so the logits it computes are the ones written here. The published headers carry shapes, not values, so whether the shipped head equals embed is not known; if it does not, a loader that uses the shipped lm_head.weight computes different logits (findings §10). granite-4.1-8b-base, whose configuration is byte-identical, ships no such tensor, which is the whole of the gap between their parameter counts (8.8B and 8.4B).

Multiplier placement follows modeling_granite.py at transformers v4.57.1: embedding, attention scale, residual, logits.

Geometry

Weight structure

The tensors one element of the repeating stack holds, by the names the published checkpoint uses.

Other shapes of GraniteForCausalLM

Identical across all of them: KV heads 8.

Only the columns that differ are shown; a cell with several values means the checkpoints of that shape disagree.

How this was checked

Status and evidence grade in the assessment (see what the grades mean):

• granite-4.1-8b, granite-4.2-8b, granite-4.1-8b-base, granite-guardian-4.1-8b, granite-3.3-8b-instruct, granite-3.3-8b-math-prm-v2, granite-guardian-3.3-8b, granite-3.0-8b-instruct, granite-3.1-8b-instruct, granite-3.2-8b-instruct, granite-3.2-8b-instruct-preview, granite-guardian-3.0-8b, granite-guardian-3.1-8b, granite-guardian-3.2-8b-factuality-detection, granite-3.0-8b-base, granite-3.1-8b-base, granite-3.3-8b-base: runs, shape-parity — a stored checkpoint of the same shape was compared against a reference
• granite-3.0-8b-lora-intrinsics-v0.1, granite-3.1-8b-lora-intrinsics-v0.1: foldable, not packed, base-shape-parity — a stored checkpoint of the compatible base was compared against a reference recording; nothing was folded for this row

Parts of the author's unpublished implementation that exist for this shape:

• a decoder implemented from the published configuration
• a configuration parser that refuses what it cannot express
• a checkpoint loader, which is where published tensor names are read
• an independent reference implementation of this architecture in Python, driving the published modelling code — an executable statement of what the model should compute, written against the publication rather than against any one implementation of it

Checkpoints with this architecture

Recorded pass: passes.