GraniteSwitchForCausalLM

GraniteSwitchForCausalLM at 41 layers and hidden size 4096. One published checkpoint has this shape.

Layer map

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

AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA

The configuration lists 41 attention layers; one of those slots is the adapter switch, which has no weights, so the checkpoint holds 40 decoder layers (see Depth). The switch is slot 0, the leftmost character: it is created with layer_idx=0 (L162), and decoder layer i takes slot i + 1 (L205–L208).

Implementing it

How to read this sheet: method.

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

GraniteSwitchForCausalLM is not in transformers, and the model repository ships no modelling file. The class is defined in generative-computing/granite-switch, which the model card names. Behaviour below is sourced to that repository at commit b261d23, the last commit before the checkpoint's current upload on 2026-05-24; tensor names are from the checkpoint's weight index.

Generated from granite-switch-4.1-8b-preview.

# The switch: no weights, one attention cache slot. For each token t,
#   c[j] = n if ids[j] is adapter n's control token (ids 100352..100363), else 0
#   a_idx[t] = round( sum over j <= t of softmax_j(+15 if c[j] else -15) * c[j] )
#   i.e. the mean adapter number of the control tokens seen so far
#   (the adapter itself when one was given); 0 before any
a_idx = switch(ids)                                  # [T], 0 = base model
ids = where(c > 0, control_to_substitute_lut[ids], ids)   # control tokens are
                                                     # embedded as a substitute
x = embed[ids] * 12                    # [T, 4096]   <- embedding_multiplier

for i in 0 .. 39:                      # 40 decoder layers; see notes
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    attn = layers[i].self_attn
    qkv = h @ attn.qkv_proj.base_layer.weight.T + lora(h, attn.qkv_proj, a_idx)
    q, k, v = split(qkv, [32*128, 8*128, 8*128])     # fused; one LoRA slice each
    q, k = rope(q, k, theta=10000000)
    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 @ attn.o_proj.base_layer.weight.T + lora(a, attn.o_proj, a_idx)
    x = x + a * 0.22        # <- residual_multiplier

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    mlp = layers[i].shared_mlp
    gu = h @ mlp.input_linear.base_layer.weight.T + lora(h, mlp.input_linear, a_idx)
    g, u = split(gu, 2)                               # [T, 12800] each; one LoRA slice each
    y = (silu(g) * u) @ mlp.output_linear.base_layer.weight.T
    y = y + lora(silu(g) * u, mlp.output_linear, a_idx)
    x = x + y * 0.22

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

# lora(z, m, a_idx), per token: zero where a_idx = 0; otherwise, with n = a_idx,
#   (z @ m.lora_A[n-1, 0].T) @ m.lora_B[n-1, 0].T
# using lora_A_slices.s / lora_B_slices.s for each output slice s of a fused
# projection. Ranks are zero-padded to 32; no scale is applied at run time.

Notes

projection_head_dim is published as 128 (4096 / 32).

Depth. num_hidden_layers is 41, but the checkpoint holds 40 decoder layers (model.layers.0 to model.layers.39). The model builds num_hidden_layers minus the switch's cache slots (L198–L199), and this switch takes one slot (single.py L107–L109).

The switch is a single attention head with no weights: query 1, key +15 at a control token and -15 elsewhere, value the adapter number, scale 1, under the causal mask; the output is rounded and clamped to 0..12 (single.py L155–L225). control_token_gain: 15 is that key magnitude. The control token itself is then replaced by its adapter_substitute_token_ids entry before embedding (L239–L244).

Adapters. 12 LoRA adapters are stored stacked in each projection, adapter n at index n-1, on qkv_proj, o_proj, shared_mlp.input_linear and shared_mlp.output_linear (lora.py L143–L159). The code applies no scale at run time: the published lora_B tensors are already multiplied by alpha/rank, which is done when the checkpoint is composed (composer/weight_transfer.py L493–L495). The attention projection is fused and split as q, k, v (L458); the attention scale is attention_multiplier (L364).

Keys the model does not use. The config carries a full mamba_* block and num_experts_per_tok: 2, keys of GraniteMoeHybridConfig, which GraniteSwitchConfig extends (config.py L9). Every layer is an attention decoder layer (L204–L208), and with num_local_experts: 0 no router or experts are built, so the feed-forward is shared_mlp alone (L53–L59, L113–L116).

Multiplier placement: embedding, residual, logits.

Geometry

Other shapes of GraniteSwitchForCausalLM

Identical across all of them: stack attention, KV heads 8, active —¹, vocabulary 100,364, context 131,072.

¹ The config sets num_experts_per_tok: 2 with num_local_experts: 0 (findings §2); with no experts, no router is built (L53–L56) and nothing is routed.

Only the columns that differ are shown; a cell with several values means the checkpoints of that shape disagree.

Why this is refused

Status and evidence grade in the assessment (see what the grades mean):

• granite-switch-4.1-8b-preview: refused by design, refusal-path-tested — this configuration is refused by name at load, with tests asserting the refusal

This configuration is refused rather than run by an unpublished independent inference implementation by the author; nothing about it is approximated.

Checkpoints with this architecture