GraniteMoeHybridForCausalLM

GraniteMoeHybridForCausalLM at 40 layers and hidden size 4096. 2 published checkpoints share this shape.

Layer map

One character per layer, input on the left, read from the published configuration.

mmmmmAmmmmmmmmmAmmmmmmmmmAmmmmmmmmmAmmmm

Implementing it

How to read this sheet: method.

Generated from granite-4.0-h-small, which covers all 2 checkpoints of this shape — they agree on every field below.

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

for i in 0 .. 39:
    # layer_types[i] decides the mixer; both share the block shape
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)

    if layer_types[i] == "mamba":
        zxbcdt = h @ layers[i].mamba.in_proj.T
        z, xBC, dt = split(zxbcdt, [8192, 8448, 128])
        xBC = silu(causal_conv1d(xBC, layers[i].mamba.conv1d, width=4, bias=layers[i].mamba.conv1d.bias))
        u, B, C = split(xBC, [8192, 128, 128])
        dt = softplus(dt + layers[i].mamba.dt_bias)
        A  = -exp(layers[i].mamba.A_log)
        dA = exp(dt * A)                   # per head, in (0,1]
        S  = dA * S + dt * outer(u, B)
        y  = S @ C + layers[i].mamba.D * u   # per token, in order
        #  ^ ssd_scan, written out: one state [64 x 128] per head,
        #    a scalar decay per head, and one B/C pair shared across
        #    every head in a group the way GQA shares a KV head.
        y  = group_rmsnorm(y * silu(z), layers[i].mamba.norm,
                           groups=1, width=8192, eps=1e-5)
        #  ^ gate first, then normalise (see notes).
        #    `mamba_chunk_size: 256` is the block length transformers
        #    uses to evaluate this same recurrence in chunks.
        y  = y @ layers[i].mamba.out_proj.T
    else:
        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]
        # no rope: position_embedding_type is "nope" (see notes)
        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 + (y if layer_types[i] == "mamba" else a) * 0.22

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    moe = layers[i].block_sparse_moe
    r = (h @ moe.router.layer.weight.T).float()     # [T, 72]
    r_top, sel = topk(r, 10)                       # tie order unspecified
    w = softmax(r_top)                             # over the 10 selected, not all 72
    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)
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)   # shared expert, width 1536
    y += (silu(g) * u) @ layers[i].shared_mlp.output_linear.T  # every token, not routed
    x = x + y * 0.22

x = rmsnorm(x, model.norm)
logits = (x @ embed.T) / 16        # tied to the input embedding; <- logits_scaling

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

State-space layer. Split sizes and order follow modeling_granitemoehybrid.py at v4.57.1: the mixer width is mamba_expand × hidden_size (L409); in_proj splits into the gate z, the convolution input and one dt per head (L651–L653); after the convolution that input splits into x, B and C (L692–L696); norm multiplies by silu(z) and then normalises, with eps = rms_norm_eps (L874–L883, L416). Which layers are state-space comes from layer_types alone (L1125–L1128); a config with no layer_types makes every layer state-space (configuration_granitemoehybrid.py L252–L253).

Multiplier placement: embedding, attention scale, residual, logits.

position_embedding_type: "nope": transformers builds no rotary embedding for this value (L1306) and the attention layers then skip rope (L197–L199). The config still carries a rope_theta, which is unused. The recorded pass logs rope because the step is entered, as annotated in its tree.

Routing. The router logits are cast to float32, the top k are taken with topk, and the gates are the softmax of those k logits only, not the top k of a softmax over all experts (L1024–L1026). PyTorch does not guarantee which index topk returns for tied values (torch.topk). transformers sorts the token–expert pairs by ascending expert index (L1040) and adds the gated expert outputs back with one index_add (L1105); summation order beyond that is not part of the model's definition. The shared MLP is added to the routed sum for every token (L1205).

Geometry

Weight structure

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

Other shapes of GraniteMoeHybridForCausalLM

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 2 checkpoints of this shape: runs, recorded-forward — this checkpoint was loaded and completed a forward pass

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

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