nnterp · encyclopedia granitemoeshared · nnterp.families.granitemoeshared

family granitemoeshared

GraniteMoE-Shared

GraniteMoE's block with a dense shared expert beside the mixture: the block adds the sum of the two times residual_multiplier, and mlp_output is that scaled sum, while mlp.output is the mixture's output alone.

Scaled residual addsMixture 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.

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

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == 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 GraniteMoeSharedForCausalLMlayoutwhat 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] GraniteMoeSharedDecoderLayerlayoutwhat 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 GraniteMoeSharedAttentionlayoutwhat 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 nn_functional_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 nn_functional_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 GraniteMoeSharedMoElayoutwhat 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: the mixture plus the shared expert, times residual_multiplier

inside the forward, at the operation that _experts_sum finds on each call
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
shared_expert_output Residual [batch seq hidden_size] The shared expert's output

the sibling module shared_mlp's output

Printout

What print(model) shows on ibm-research/moe-7b-1b-active-shared-experts: the native modules, standard names and every standard value with its layout.

GraniteMoeSharedForCausalLM(
  (lm_head): Linear(in_features=1536, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50257, 1536, padding_idx=0)
  (layers): ModuleList(
    (0-39): 40 x GraniteMoeSharedDecoderLayer(
      (self_attn): GraniteMoeSharedAttention(
        (q_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (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_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (input_layernorm): GraniteMoeSharedRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeSharedRMSNorm((1536,), eps=1e-05)
      (mlp/block_sparse_moe): GraniteMoeSharedMoE(
        (router): GraniteMoeSharedTopKRouter()
        (experts): GraniteMoeSharedExperts(
          (act_fn): SiLUActivation()
        )
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the mixture plus the shared expert, times residual_multiplier
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (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_mlp): GraniteMoeSharedMLP(
        (activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (output_linear): Linear(in_features=1024, out_features=1536, 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): GraniteMoeSharedRMSNorm((1536,), 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_layers40
model.hidden_size1536
model.num_heads12
model.num_kv_heads4
model.head_dim128
model.qk_head_dim128
model.vocab_size50257
model.intermediate_size512
architecturesGraniteMoeSharedForCausalLM
num_hidden_layers40
hidden_size1536
intermediate_size512
num_attention_heads12
num_key_value_heads4
vocab_size50257
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
embedding_multiplier12
residual_multiplier0.22
attention_multiplier0.0078125
logits_scaling6
num_experts_per_tok6
num_local_experts62
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

GraniteMoE's block with a dense expert beside the mixture on every block. The mixture's native name, block_sparse_moe, is mlp; the shared expert keeps its native name, shared_mlp, on the block. The released checkpoint, moe-7b-1b-active-shared-experts, has 40 blocks, each with 62 experts 512 wide, the top 6 routed per token, and a shared expert 1024 wide (shared_intermediate_size), a SwiGLU like the experts. Its multipliers are Granite's: embedding_multiplier 12, residual_multiplier 0.22, attention_multiplier 1/128 and logits_scaling 6.

mlp_output is both experts' sum, scaled

No module returns what the block adds after the attention, so mlp_output is the block's own binding of block_sparse_moe(n) + shared_mlp(n), times residual_multiplier. The mixture's values are read before the multiplier: routed_output is mlp.output, the mixture's output alone, and shared_expert_output is layers[i].shared_mlp.output. The routed experts run first, so read routed_output before shared_expert_output, and both before mlp_output:

moe = model.layers[1].mlp
with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    mlp = moe.mlp_output.save()
    out = model.layers[1].layer_output.save()

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

Attribute to single experts with expert_outputs * m and to the shared expert with shared_expert_output * m: those are their terms in the stream.

A write to mlp_output replaces both experts

A write to mlp_output reaches the model on the sum: the block's binding is replaced by the edited copy divided by residual_multiplier, so mlp_output[:] = 0 removes the mixture and the shared expert together and leaves layer_output equal to input + attention_output. To ablate one of the two, edit its module's output: mlp.output for the mixture, layers[i].shared_mlp.output for the shared expert, scaled by 1 / residual_multiplier to land as a given vector on the stream.

m = model.config.residual_multiplier
with model.trace(prompt):
    model.layers[1].shared_mlp.output[:, -1] += v / m    # the shared expert's term moves by v

Like every Granite write to a scaled copy, this divides the whole copy back, so positions you did not edit are rewritten too. With 0.22 the round trip is not exact in bfloat16: for 502 of the 65280 finite products x * 0.22, / 0.22 * 0.22 does not give the product back. Load in float32 for fine-grained edits.

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

GraniteMoE's routing: one linear map without a bias, the top num_experts_per_tok of its logits, and a softmax over those alone, so a token's expert_weights sum to one. expert_outputs (each slot's weighted output) sum to routed_output. The shared expert has no weight: it is added to every token as it is.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()        # [batch, seq, num_experts]
    w = moe.expert_weights.save()            # [batch, seq, top_k]
    idx = moe.expert_indices.save()
    slots = moe.expert_outputs.save()
    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)

intermediate_size and mlp.intermediate_size are the routed experts' width, 512; the shared expert's is config.shared_intermediate_size. A config with shared_intermediate_size 0 builds no shared expert: shared_mlp is None, shared_expert_output is unavailable, and mlp_output is the mixture's output times the multiplier.

Attention, logits and embeddings are Granite's

The attention scales q · kᵀ by attention_multiplier, 1/128, which is 1 / head_dim. There are 12 query heads over 4 key/value heads of 128, so an edit to key/value head j reaches query heads 3j to 3j + 2. The six interior values need attn_implementation="eager".

model.logits is model.lm_head.output / logits_scaling and project_on_vocab divides too; token_embeddings is the plain lookup and layers[0].input is it times 12. 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 tokenizer is Granite 3's, 49152 tokens (the embedding has 50257 rows), and prepends nothing. It splits Paris into ĠPar and is; London is one token.

What this family module covers

Every checkpoint whose config says granitemoeshared: ibm-research's moe-7b-1b-active-shared-experts. granitemoe is the same block without the shared expert; granitemoehybrid (Granite 4.0) has the same feed-forward, with Mamba-2 mixers on most blocks; granitemoe_swa has a shared expert too, which its mlp_output leaves out.

From the family module

04 · nnterp/families/granitemoeshared.py

GraniteMoE with a shared expert (GraniteMoeSharedForCausalLM).

GraniteMoE's block (granitemoe.py) with a dense shared expert beside the mixture: the block computes block_sparse_moe(h) + shared_mlp(h) from the pre-MLP norm's output h and adds that sum times residual_multiplier. No module returns the sum, so block_sparse_moe is mlp (it is on every block; shared_mlp is None when shared_intermediate_size is 0) and its mlp_output is the block's binding of the sum (hidden_states_4, or hidden_states_3, the mixture's output alone, without a shared expert) times the multiplier: what the block adds. As on Granite it is a computed copy, divided on assignment and carried back by a transform after an in-place edit; the write lands on the sum, so it replaces both experts' contribution. mlp.output is the mixture's output alone and layers[i].shared_mlp.output the shared expert's. The binding is in the block's forward, read after the block has started, so the family's Layer sets sourced; the mixture keeps no config, so its Mlp reads the multiplier and whether the block has a shared expert off its parent block (residual_multiplier, shared). The attention, embedding_multiplier and logits_scaling are as on Granite. The mixture is a Moe (GraniteMoE's routing); its shared_expert_output is the block's shared_mlp's output, unscaled, so routed_output + shared_expert_output == mlp_output / residual_multiplier.