GraniteMoeHybridForCausalLM

GraniteMoeHybridForCausalLM at 40 layers and hidden size 2048. 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.

No forward was recorded for this shape; the block is read from the published configuration.

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

x = embed[ids] * 12                    # [T, 2048]   <- 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, [4096, 4352, 64])
        xBC = silu(causal_conv1d(xBC, layers[i].mamba.conv1d, width=4, bias=layers[i].mamba.conv1d.bias))
        u, B, C = split(xBC, [4096, 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=4096, 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*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]
        # 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.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.22

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)   # width 8192
    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.22

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

head_dim is not published; 2048 / 32 = 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.

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

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, 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 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