GraniteForCausalLM

GraniteForCausalLM at 40 layers and hidden size 2560. 3 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From `granite-4.1-3b`

Also covers granite-4.1-3b-base.

x = embed[ids] * 12                    # [T, 2560]   <- 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, 40*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=10000000)
    k, v = repeat_kv(k, v, 5)            # 40 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) / 10        # tied to the input embedding; <- logits_scaling

Differs from the block above

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

checkpoint	`embedding_multiplier`	`residual_multiplier`	`logits_scaling`	output head
`granite-4.2-3b`	1	1	1	untied: `logits = (x @ lm_head.T) / 1`

Notes

head_dim is not published; 2560 / 40 = 64 is used.

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

Geometry


layers	40
hidden size	2,560
attention heads	40, 8 key/value
feed-forward width	8,192
vocabulary	100,352
trained context	131,072 tokens
largest checkpoint	3.7B

Weight structure

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

repeating stack	depth	what one element holds
`model.layers.#`	40	`input_layernorm.weight`, `mlp.down_proj.weight`, `mlp.gate_proj.weight`, `mlp.up_proj.weight`, `post_attention_layernorm.weight`, `self_attn.k_proj.weight`, `self_attn.o_proj.weight`, `self_attn.q_proj.weight`, `self_attn.v_proj.weight`

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.

shape	layers	width	heads	FFN width	vocabulary	context	checkpoints
`64L x 4,096`	64	4,096	32	32,768	100,352	131,072	3
`40L x 4,096`	40	4,096	32	12,800	49,152 / 49,155 / 49,159 / 100,352	4,096 / 8,192 / 131,072	19
`40L x 2,560` (this sheet)	40	2,560	40	8,192	100,352	131,072	3
`40L x 2,048`	40	2,048	32	8,192	49,152 / 49,155 / 49,159	4,096 / 8,192 / 131,072	9
`28L x 4,096`	28	4,096	32	12,800	49,155	131,072	2

How this was checked

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

each of the 3 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

model	parameters	context
`granite-4.2-3b`	3.7B	131,072
`granite-4.1-3b`	3.4B	131,072
`granite-4.1-3b-base`	3.4B	131,072

Recorded pass: passes.