GraniteMoeHybridForCausalLM

GraniteMoeHybridForCausalLM at 40 layers and hidden size 2560. 3 published checkpoints share this shape.

Layer map

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

AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA

Implementing it

How to read this sheet: method.

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

# The config also carries `mamba_*` keys (d_state 256, 128 heads,
#   1 group, conv width 4, chunk 256). `layer_types` lists
#   every layer as attention and transformers picks each layer's
#   mixer from that list alone, so the keys build nothing here.

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

for i in 0 .. 39:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T       # [T, 40*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, 5)            # 40 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)
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)   # [T, 8192] each
    y = (silu(g) * u) @ layers[i].shared_mlp.output_linear.T
    x = x + y * 0.22

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

head_dim is not published; 2560 / 40 = 64 is used.

Layer kinds. transformers builds each layer's mixer from layer_types alone (L1125–L1128), so with every entry attention the mamba_* keys build nothing. Without layer_types it would make every layer state-space (configuration_granitemoehybrid.py L252–L253).

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

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 3 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 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.