LlamaForCausalLM

LlamaForCausalLM at 36 layers and hidden size 4096. 4 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From granite-8b-code-base-128k

Also covers granite-8b-code-instruct-128k.

x = embed[ids]                    # [T, 4096]

for i in 0 .. 35:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T + bias       # [T, 32*128]
    k = h @ layers[i].self_attn.k_proj.T + bias       # [T, 8*128]
    v = h @ layers[i].self_attn.v_proj.T + bias       # [T, 8*128]
    # `attention_bias: true` puts a bias on all four attention
    #   projections, o_proj included (see Notes).
    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 / sqrt(128), mask=causal) @ v
    a = a @ layers[i].self_attn.o_proj.T + bias
    x = x + a

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g = silu(h @ layers[i].mlp.gate_proj.T + ...bias)           # [T, 14336]
    u =       h @ layers[i].mlp.up_proj.T + ...bias
    y = (g * u) @ layers[i].mlp.down_proj.T + ...bias
    # `mlp_bias: true` — three more bias tensors, one per projection.
    x = x + y

x = rmsnorm(x, model.norm)
logits = x @ embed.T        # tied to the input embedding

Differs from the block above

The same arithmetic with this value changed. Everything not listed is as above.

Notes

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

In transformers' LlamaAttention at v4.57.1, attention_bias sets the bias on q, k, v and o_proj (L210-L221), and mlp_bias on all three MLP projections (L149-L151); the checkpoint ships all seven bias tensors. Qwen2 has no such flag: it hard-codes a bias on q, k and v and none on o_proj (modeling_qwen2.py L134-L137), so a port that borrows Qwen2's attention drops o_proj.bias.

Geometry

Weight structure

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

Other shapes of LlamaForCausalLM

Identical across all of them: heads 32.

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 4 checkpoints of this shape: runs, class — the architecture class is implemented; nothing checkpoint-specific was measured

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

• a decoder implemented from the published configuration
• 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.