nnterp · encyclopedia granitemoe_swa · nnterp.families.granitemoe_swa

family granitemoe_swa

Granite Swash MoE

Granite Swash's sliding-window attention with its sink in GraniteMoE's scaled block, where a shared expert runs beside the mixture of experts and mlp_output carries the mixture's term alone.

Scaled residual addsSliding windowSink-scaled headsMixture of expertsEmbedding multiplierScaled logitsGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 28
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsscaled
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output + shared_mlp.output * residual_multiplier == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model GraniteMoeSWAForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] GraniteMoeSWADecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn GraniteMoeSWAAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of F_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream: its output times residual_multiplier

the module's own output
attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of F_dropout_0 inside attention_interface_1
attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp GraniteMoeSWAMoElayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream: its output times residual_multiplier

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

Printout

What print(model) shows on ibm-granite/granite-swash-3b-a600m: the native modules, standard names and every standard value with its layout.

GraniteMoeSWAForCausalLM(
  (lm_head): Linear(in_features=1280, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 1280, padding_idx=100256)
  (layers): ModuleList(
    (0-27): 28 x GraniteMoeSWADecoderLayer(
      (self_attn): GraniteMoeSWAAttention(
        (q_proj): Linear(in_features=1280, out_features=1280, bias=False)
        (k_proj): Linear(in_features=1280, out_features=256, bias=False)
        (v_proj): Linear(in_features=1280, out_features=256, bias=False)
        (o_proj): Linear(in_features=1280, out_features=1280, bias=False)
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (input_layernorm): GraniteMoeSWARMSNorm((1280,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeSWARMSNorm((1280,), eps=1e-05)
      (mlp/block_sparse_moe): GraniteMoeSWAMoE(
        (router): GraniteMoeSWATopKRouter()
        (experts): GraniteMoeSWAExperts(
          (act_fn): SiLUActivation()
        )
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: its output times residual_multiplier
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
      )
      (shared_mlp): GraniteMoeSWAMLP(
        (activation): SiLUActivation()
        (input_linear): Linear(in_features=1280, out_features=2560, bias=False)
        (output_linear): Linear(in_features=1280, out_features=1280, bias=False)
      )
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): GraniteMoeSWARMSNorm((1280,), eps=1e-05)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers28
model.hidden_size1280
model.num_heads20
model.num_kv_heads4
model.head_dim64
model.qk_head_dim64
model.vocab_size100352
model.intermediate_size512
architecturesGraniteMoeSWAForCausalLM
num_hidden_layers28
hidden_size1280
intermediate_size512
num_attention_heads20
num_key_value_heads4
vocab_size100352
max_position_embeddings8192
sliding_window128
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
embedding_multiplier12
residual_multiplier0.26
attention_multiplier0.015625
logits_scaling5
num_experts_per_tok4
num_local_experts48
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

x      = embed_tokens(ids) * embedding_multiplier
h      = x + self_attn(input_layernorm(x)) * residual_multiplier
n      = post_attention_layernorm(h)
out    = h + (block_sparse_moe(n) + shared_mlp(n)) * residual_multiplier
logits = lm_head(norm(out)) / logits_scaling

The attention is Granite Swash's: of granite-swash-3b-a600m's 28 blocks, 8 attend over the whole prefix (0, 3, 7, 11, 15, 19, 23 and 27) and 20 over the last 128 tokens, every block applies rotary embeddings with base 10000, and each head carries a sink that scales its output. The feed-forward is a mixture of 48 experts, 512 wide, with the top 4 routed per token, and beside it shared_mlp, a SwiGLU 1280 wide that runs on every token. The multipliers: embedding_multiplier 12, residual_multiplier 0.26, attention_multiplier 1/64 and logits_scaling 5.

mlp_output leaves the shared expert out

The mixture's native name, block_sparse_moe, is mlp, and mlp_output is its output times residual_multiplier. The block adds the shared expert's output times the multiplier too, and no standard value carries that term: shared_expert_output is unavailable ("this mixture has no shared expert"), and layers[i].shared_mlp keeps its native name. The block's sum therefore has a third term:

layer = model.layers[1]
with model.trace(prompt):
    x = layer.input.save()
    attn = layer.self_attn.attention_output.save()
    mlp = layer.mlp.mlp_output.save()
    shared = layer.shared_mlp.output.save()
    out = layer.layer_output.save()

m = model.config.residual_multiplier                 # 0.26
torch.testing.assert_close(x + attn + mlp + shared * m, out)

The term is not small. Over the positions after the first, the shared expert's scaled output has a mean norm of 3.5 in block 1 against 3.4 for mlp_output, and 2.4 against 2.0 in block 14; x + attn + mlp misses layer_output by up to 0.53 in block 1 and 8.1 in block 14. Decompose the stream, attribute or ablate the feed-forward with both terms: a write to mlp_output replaces the mixture's term and leaves the shared expert's, so zeroing the whole feed-forward means zeroing mlp_output and shared_mlp.output.

The router: top 4 of the logits, then a softmax over those 4

The router is one linear map without a bias. It takes the top 4 of its 48 logits and softmaxes those 4 alone, so a token's expert_weights sum to one, and expert_outputs (each slot's weighted expert output) sum to routed_output, which is mlp.output, unscaled:

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()        # [batch, seq, 48]
    w = moe.expert_weights.save()            # [batch, seq, 4]
    idx = moe.expert_indices.save()
    slots = moe.expert_outputs.save()        # [batch, seq, 4, hidden_size]
    routed = moe.routed_output.save()

top, chosen = logits.topk(moe.top_k, dim=-1)
torch.testing.assert_close(chosen, idx)
torch.testing.assert_close(top.float().softmax(-1).to(w.dtype), w)
torch.testing.assert_close(slots.sum(2), routed)

The experts are stored fused: experts.gate_up_proj is [48, 1024, 1280], each expert's gate rows then its up rows, and experts.down_proj is [48, 1280, 512]. The default experts_implementation, grouped_mm, serves expert_outputs.

The sink scales the head outputs

As on Granite Swash, each head's output is multiplied after the softmax by sigmoid(logsumexp(scores) - sink), with the sink in self_attn._module.sinks, so attention_probabilities rows sum to one and attention_head_outputs are read after the scale. On a code prompt the median of that share over heads and queries is 0.84 in block 0, and 0.14 to 0.25 in blocks 1, 3 and 27. A pattern read as what a head moves is probs * share[..., None]; to set what a head writes, edit attention_head_outputs.

attn = model.layers[1].self_attn
with model.trace(prompt):
    scores = attn.attention_scores.save()
    probs = attn.attention_probabilities.save()

share = torch.sigmoid(scores.logsumexp(-1) - attn._module.sinks.view(1, -1, 1))
torch.testing.assert_close(scores.softmax(-1), probs)

transformers has no sdpa path for this attention, so a default load runs eager and the six interior values are available. The query scale is attention_multiplier, 1/64, which is 1 / head_dim. There are 20 query heads over 4 key/value heads of 64: an edit to key/value head j reaches query heads 5j to 5j + 4.

A write moves the other positions in bfloat16

A write to attention_output or mlp_output reaches the model as the whole edited copy divided by residual_multiplier. With 0.26 that round trip is exact for all but 2 of the 65280 finite bfloat16 products x * 0.26, but the mixtures add a second source: an edit changes the edited token's routing in later blocks, and from then on the grouped expert kernels compute the other tokens' outputs with different rounding. Adding a unit vector at the last position of block 1, to mlp_output or as mlp.output[:, -1] += v / m, moved the earlier positions' logits by up to 0.62 (a KL of 0.004) in bfloat16, against a KL of 9e-6 at the edited position. In float32 the earlier positions stay bit-identical. Load in float32 for any edit whose effect is small.

m = model.config.residual_multiplier
with model.trace(prompt):
    model.layers[1].mlp.output[:, -1] += v / m       # bfloat16: earlier positions move by rounding

Logits are divided by 5, and the embeddings multiplied by 12

model.logits is model.lm_head.output / logits_scaling exactly, and project_on_vocab divides too. A softmax over lm_head.output is at temperature 1/5: after "The capital of the United Kingdom is", London has probability 0.70 from logits and 1.00 from lm_head.output. token_embeddings is the plain lookup and layers[0].input is it times 12; to set what block 0 reads, edit layers[0].input. lm_head and embed_tokens are one parameter.

with model.trace(prompt):
    emb = model.token_embeddings.save()
    first = model.layers[0].input.save()
    resid = model.layers[-1].layer_output.save()
    raw = model.lm_head.output.save()
    logits = model.logits.save()

torch.testing.assert_close(first, emb * model.config.embedding_multiplier)
torch.testing.assert_close(logits, raw / model.config.logits_scaling)
torch.testing.assert_close(model.project_on_vocab(resid), logits)

The vocabulary is Granite 4's, 100352 tokens; no beginning-of-sequence token is prepended, and " London" (7295) and " Paris" (12366) are single tokens.

What this family module covers

Every checkpoint whose config says granitemoe_swa: granite-swash-3b-a600m. granite_swa is the dense Swash model, the same attention with an MLP. granitemoeshared is GraniteMoE's block with a shared expert, whose mlp_output carries both experts' sum.

From the family module

04 · nnterp/families/granitemoe_swa.py

GraniteMoE SWA (GraniteMoeSWAForCausalLM).

Granite SWA's attention (granite_swa.py: sliding-window and full blocks, a per-head sink that scales the head outputs after the softmax) in GraniteMoE's block (granitemoe.py): block_sparse_moe, routed experts only, is mlp, and the block adds each sublayer's output times residual_multiplier, so attention_output and mlp_output are the scaled copies. embedding_multiplier and logits_scaling are as on Granite.

Where the config sets shared_intermediate_size (1280 on ibm-granite/granite-swash-3b-a600m; 0 on the pinned tiny) the block also runs a shared_mlp beside the experts and adds the sum. This family binds mlp_output to the routed experts' output times residual_multiplier only, and shared_expert_output reports no shared expert, so on that checkpoint input + attention_output + mlp_output misses layer_output by the shared expert's term; read layers[i]._module.shared_mlp (or its envoy) for it.