GraniteMoeForCausalLM

GraniteMoeForCausalLM at 32 layers and hidden size 1536. 5 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

From granite-3.1-3b-a800m-instruct

Also covers granite-guardian-3.2-3b-a800m, granite-3.1-3b-a800m-base.

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

for i in 0 .. 31:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-06)
    q = h @ layers[i].self_attn.q_proj.T       # [T, 24*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, 3)            # 24 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)
    moe = layers[i].block_sparse_moe
    r = (h @ moe.router.layer.weight.T).float()     # [T, 40]
    r_top, sel = topk(r, 8)                        # tie order unspecified
    w = softmax(r_top)                             # over the 8 selected, not all 40
    y = 0
    for (e, w_e) in zip(sel, w):                   # summation order: see notes
        g, u = split(h @ moe.input_linear.weight[e].T, 2)
        y += w_e * ((silu(g) * u) @ moe.output_linear.weight[e].T)
    x = x + y * 0.22

x = rmsnorm(x, model.norm)
logits = (x @ embed.T) / 6        # 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; 1536 / 24 = 64 is used.

Output head: the configuration says tied, the checkpoints ship a separate one. granite-3.0-3b-a800m-base and granite-3.0-3b-a800m-instruct set tie_word_embeddings: true, which is what the block follows, but their checkpoints also ship a separate lm_head.weight ([49152, 1536] and [49155, 1536]). 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).

Routing. The router logits are cast to float32, the top 8 are taken with topk, and the gates are the softmax of those 8 logits only, not the top 8 of a softmax over all experts (modeling_granitemoe.py L294–L296). PyTorch does not guarantee which index topk returns for tied values (torch.topk). transformers sorts the token–expert pairs by ascending expert index (L310) and adds the gated expert outputs back with one index_add (L371); summation order beyond that is not part of the model's definition.

Multiplier placement: 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 GraniteMoeForCausalLM

Identical across all of them: KV heads 8, FFN width 512, active 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 5 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 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.