Safe By Design AI

LlamaForCausalLM

LlamaForCausalLM at 32 layers and hidden size 2560. 4 published checkpoints share this shape.

The pass (full forward)

The order of operations for one forward. ×N marks a position that fires once per layer.

17 steps across nested levels, recorded from one complete forward of this shape in the author's independent implementation — the structure is what that run emitted, not a reading of the configuration.

forward
├─ embed
├─ layer ×N
│  ├─ rmsnorm_attn
│  ├─ wq
│  ├─ wk
│  ├─ wv
│  ├─ rope
│  ├─ attn
│  │  ├─ scores
│  │  └─ attn_mix
│  ├─ wo
│  ├─ rmsnorm_mlp
│  ├─ w_gate
│  ├─ w_up
│  └─ w_down
└─ head

Recorded separately from the structure: layer fired 32 times. A ×N run says only that it repeated — how long a run is is not part of the structure, so two models differing only in depth have the same pass.

The same thing in canonical form:

0(1 2*N(3 4 5 6 7 8(9 10) 11 12 13 14 15) 16)

Numbers are positions in the pass, not layer indices. Two passes count as the same structure when these strings match.

Implementing it

One forward, written out, per checkpoint whose arithmetic actually differs. A shape groups checkpoints by class, depth, width and layer stack — none of which decides an activation, a rope base or a routing rule — so where the members of this shape disagree there is a block each, and each names the checkpoint it was generated from.

It is not the recorded pass above, which is captured from a real forward; it is what the published configuration says the arithmetic is. Where the configuration does not say, the line says that instead of guessing.

The names are the ones the published checkpoint uses where a name is published, and canonical otherwise. A checkpoint loader may store them differently — fusing a gate/up pair into one matrix, or splitting one published projection in two — and those are that loader's names, not the model's. Everything spelled here is what the download contains.

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

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

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

for i in 0 .. 31:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T + bias       # [T, 32*80]
    k = h @ layers[i].self_attn.k_proj.T + bias       # [T, 32*80]
    v = h @ layers[i].self_attn.v_proj.T + bias       # [T, 32*80]
    # Biases are read from the checkpoint, not from the flag.
    #   `attention_bias: true` covers q/k/v; `o_proj.bias` is a
    #   separate tensor and this checkpoint ships it. Qwen2 sets the
    #   flag and ships no o_proj bias, so one bit cannot describe the
    #   set — a loader has to look for each of the four.
    q, k = rope(q, k, theta=10000000)
    a = softmax(q @ k.T / sqrt(80), 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, 10240]
    u =       h @ layers[i].mlp.up_proj.T + ...bias
    y = (g * u) @ layers[i].mlp.down_proj.T + ...bias
    # `mlp_bias: true` — three more tensors, and on granite-8b-code
    # `mlp.gate_proj.bias` has rms 1.88. Dropping it is not small.
    x = x + y

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

From granite-3b-code-base-2k

Also covers granite-3b-code-instruct-2k.

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

for i in 0 .. 31:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T + bias       # [T, 32*80]
    k = h @ layers[i].self_attn.k_proj.T + bias       # [T, 32*80]
    v = h @ layers[i].self_attn.v_proj.T + bias       # [T, 32*80]
    # Biases are read from the checkpoint, not from the flag.
    #   `attention_bias: true` covers q/k/v; `o_proj.bias` is a
    #   separate tensor and this checkpoint ships it. Qwen2 sets the
    #   flag and ships no o_proj bias, so one bit cannot describe the
    #   set — a loader has to look for each of the four.
    q, k = rope(q, k, theta=10000)
    a = softmax(q @ k.T / sqrt(80), 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, 10240]
    u =       h @ layers[i].mlp.up_proj.T + ...bias
    y = (g * u) @ layers[i].mlp.down_proj.T + ...bias
    # `mlp_bias: true` — three more tensors, and on granite-8b-code
    # `mlp.gate_proj.bias` has rms 1.88. Dropping it is not small.
    x = x + y

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

Notes that apply to more than one block

Stated once here rather than under each block above.

head_dim is not published; 2560 / 32 = 80 is used.

Geometry

layers32
hidden size2,560
attention heads32
feed-forward width10,240
vocabulary49,152
trained context128,000 tokens (granite-3b-code-base-128k, granite-3b-code-instruct-128k); 2,048 tokens (granite-3b-code-base-2k, granite-3b-code-instruct-2k)
largest checkpoint3.5B

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.#32input_layernorm.weight, mlp.down_proj.bias, mlp.down_proj.weight, mlp.gate_proj.bias, mlp.gate_proj.weight, mlp.up_proj.bias, mlp.up_proj.weight, post_attention_layernorm.weight, self_attn.k_proj.bias, self_attn.k_proj.weight, self_attn.o_proj.bias, self_attn.o_proj.weight, self_attn.q_proj.bias, self_attn.q_proj.weight, self_attn.v_proj.bias, self_attn.v_proj.weight

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.

shapelayerswidthKV headsFFN widthvocabularycontextcheckpoints
36L x 4,096364,096814,33649,1524,096 / 128,0004
32L x 4,096324,0963210,928 / 11,00832,000 / 32,008 / 49,1524,096 / 8,19211
32L x 2,560 (this sheet)322,5603210,24049,1522,048 / 128,0004

How this was checked

implemented, evidence grade class, per the assessment, for each checkpoint of this shape: the architecture class is implemented; nothing specific to this checkpoint was measured. The implementation meant here and below is an unpublished independent inference implementation by the author.

What stands behind the block above, beyond the published configuration it is read from:

Where an independent implementation and the published configuration disagree, the configuration is what this page reports and the disagreement is what it says.

Checkpoints with this architecture

modelparameterscontext
granite-3b-code-base-128k3.5B128,000
granite-3b-code-base-2k3.5B2,048
granite-3b-code-instruct-128k3.5B128,000
granite-3b-code-instruct-2k3.5B2,048