GraniteForCausalLM

GraniteForCausalLM at 40 layers and hidden size 2048. 9 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From granite-3.1-2b-instruct

Also covers granite-3.2-2b-instruct, granite-guardian-3.1-2b, granite-3.1-2b-base.

x = embed[ids] * 12                    # [T, 2048]   <- 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*64]
    k = h @ layers[i].self_attn.k_proj.T       # [T, 8*64]
    v = h @ layers[i].self_attn.v_proj.T       # [T, 8*64]
    q, k = rope(q, k, theta=5000000)
    k, v = repeat_kv(k, v, 4)            # 32 query heads share 8 key/value heads
    a = softmax(q @ k.T * 0.015625, 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, 8192]
    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) / 8        # tied to the input embedding; <- logits_scaling

Differs from the block above

Only the rope base differs. Everything else is as above.

Notes

head_dim is not published; 2048 / 32 = 64 is used.

Output head: the configuration says tied, the checkpoint ships a separate one. granite-3.0-2b-instruct sets tie_word_embeddings: true, which is what the block follows, but its checkpoint also ships a separate lm_head.weight [49155, 2048]. 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).

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):

• each of the 9 checkpoints of this shape: runs, shape-parity — a stored checkpoint of the same shape was compared against a reference

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.