GraniteMoeHybridForCausalLM

GraniteMoeHybridForCausalLM at 32 layers and hidden size 768. 2 published checkpoints share this shape.

Layer map

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

mmmmmmmmmmAmmAmmmAmmmmmmmmmAmmmm

Implementing it

How to read this sheet: method.

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

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

for i in 0 .. 31:
    # 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, [1536, 1792, 48])
        xBC = silu(causal_conv1d(xBC, layers[i].mamba.conv1d, width=4, bias=layers[i].mamba.conv1d.bias))
        u, B, C = split(xBC, [1536, 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 [32 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=1536, 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, 12*64]
        k = h @ layers[i].self_attn.k_proj.T       # [T, 4*64]
        v = h @ layers[i].self_attn.v_proj.T       # [T, 4*64]
        # no rope: position_embedding_type is "nope" (see notes)
        k, v = repeat_kv(k, v, 3)            # 12 query heads share 4 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 + (y if layer_types[i] == "mamba" else a) * 0.246

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)   # width 2048
    y = (silu(g) * u) @ layers[i].shared_mlp.output_linear.T
    # num_local_experts: 0, so no router and no experts (see notes)
    x = x + y * 0.246

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

head_dim is not published; 768 / 12 = 64 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.

num_local_experts: 0: transformers builds no block_sparse_moe for such a layer (L1116–L1117) and the feed-forward is shared_mlp alone (L1202–L1207).

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.