Safe By Design AI

GraniteForCausalLM

GraniteForCausalLM at 64 layers and hidden size 4096. 3 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From granite-4.1-30b

Also covers granite-4.1-30b-base.

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

for i in 0 .. 63:
    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=50000000)
    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.175        # <- residual_multiplier

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

x = rmsnorm(x, model.norm)
logits = (x @ embed.T) / 16        # 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.

checkpointembedding_multiplierresidual_multiplierlogits_scalingoutput head
granite-4.2-30b111untied: logits = (x @ lm_head.T) / 1

Notes

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

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

Geometry

layers64
hidden size4,096
attention heads32, 8 key/value
feed-forward width32,768
vocabulary100,352
trained context131,072 tokens
largest checkpoint29.3B

Weight structure

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

repeating stackdepthwhat one element holds
model.layers.#64input_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.

shapelayerswidthheadsFFN widthvocabularycontextcheckpoints
64L x 4,096 (this sheet)644,0963232,768100,352131,0723
40L x 4,096404,0963212,80049,152 / 49,155 / 49,159 / 100,3524,096 / 8,192 / 131,07219
40L x 2,560402,560408,192100,352131,0723
40L x 2,048402,048328,19249,152 / 49,155 / 49,1594,096 / 8,192 / 131,0729
28L x 4,096284,0963212,80049,155131,0722

How this was checked

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

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

Checkpoints with this architecture

modelparameterscontext
granite-4.2-30b29.3B131,072
granite-4.1-30b28.9B131,072
granite-4.1-30b-base28.9B131,072

Recorded pass: passes.