nnterp · encyclopedia jetmoe · nnterp.families.jetmoe

family jetmoe

JetMoE

Llama's pre-norm block in which both sublayers are mixtures of experts: besides the MLP's, a router picks attention experts per token, each with its own query and output projections over shared keys and values.

Attention expertsMixture of expertsGrouped-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× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
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 JetMoeForCausalLMlayoutwhat 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] JetMoeDecoderLayerlayoutwhat 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].mlp JetMoeMoElayoutwhat 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 module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

the child module router.layer's output
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

inside router's forward: the output of top_k_gates_0
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

inside router's forward: the output of logits_topk_0, element 1
⚠︎expert_outputs scalar Unavailable: JetMoE's mixture runs its experts over the slots sorted by expert; no tensor holds them in token order

nowhere: the family does not serve it
routed_output Residual [batch seq hidden_size] The routed experts' combined output, before the mixture's bias

inside the forward: the output of layer_output_1
model.layers[i].self_attn JetMoeAttentionlayoutwhat 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 scalar Unavailable: a mixture of attention heads: the output projection is per expert, inside experts

nowhere: the family does not serve it
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

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

Printout

What print(model) shows on jetmoe/jetmoe-8b: the native modules, standard names and every standard value with its layout.

JetMoeForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=32000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(32000, 2048)
  (layers): ModuleList(
    (0-23): 24 x JetMoeDecoderLayer(
      (mlp): JetMoeMoE(
        (act_fn/activation): SiLUActivation()
        (input_linear): JetMoeParallelExperts()
        (down_proj/output_linear): JetMoeParallelExperts()
        (router): JetMoeTopKGating(
          (layer): Linear(in_features=2048, out_features=8, bias=False)
        )
        (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): Unavailable: JetMoE's mixture runs its experts over the slots sorted by expert; no tensor holds them in token order
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, before the mixture's bias
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JetMoeRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): JetMoeRMSNorm((2048,), eps=1e-06)
      (self_attn/self_attention): JetMoeAttention(
        (experts): JetMoeMoA(
          (input_linear): JetMoeParallelExperts()
          (down_proj/output_linear): JetMoeParallelExperts()
          (router): JetMoeTopKGating(
            (layer): Linear(in_features=2048, out_features=8, bias=False)
          )
        )
        (kv_proj): Linear(in_features=2048, out_features=4096, bias=False)
        (attention_premix): Unavailable: a mixture of attention heads: the output projection is per expert, inside experts
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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
      )
      (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): JetMoeRMSNorm((2048,), 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_layers24
model.hidden_size2048
model.num_heads32
model.num_kv_heads16
model.head_dim128
model.qk_head_dim128
model.vocab_size32000
model.intermediate_size5632
architecturesJetMoEForCausalLM
num_hidden_layers24
hidden_size2048
intermediate_size5632
num_attention_heads32
num_key_value_heads16
vocab_size32000
max_position_embeddings4096
rms_norm_eps1e-05
tie_word_embeddingsTrue
n_positions4096
activation_functionsilu
layer_norm_epsilon1e-05
num_experts_per_tok2
num_local_experts8
kv_channels128

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attention(input_layernorm(x))   # 2 of 8 attention experts per token
out = h + mlp(post_attention_layernorm(h))     # 2 of 8 MLP experts per token, + bias

Llama's pre-norm block, with the attention under self_attention, aliased self_attn. Both sublayers are mixtures with routers of their own. The attention returns (attn_output, attn_weights, router_logits) and attention_output is the first element; the block returns a tensor, and the identity is the plain sum:

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

torch.testing.assert_close(x + attn + mlp, out)

The attention is a mixture of experts

self_attn.experts holds 8 attention experts, each a query projection (experts.input_linear, hidden_size → num_kv_heads × head_dim) and an output projection (experts.output_linear, back to hidden_size); there is no q_proj or o_proj. Per token, experts.router takes the top 2 of 8 logits and a softmax over those 2. Each chosen expert projects the token's queries for 16 heads, so a token has 2 × 16 = 32 query heads (num_heads), and slot s holds heads 16 s to 16 s + 15. Keys and values come from one shared kv_proj, 16 heads of 128, and are tiled (repeat, not interleaved) to all 32 heads: head h is slot h // 16 over key/value head h % 16.

attention_output is the gated sum of the chosen experts' output projections, each over its slot's 16 head outputs, plus the mixture's bias experts.bias. nnterp serves no standard value for the attention's routing. Its logits are experts.router.layer.output, flat over tokens ([batch × seq, num_local_experts]); the third element of self_attn.output holds them in float32:

attn = model.layers[1].self_attn
with model.trace(prompt):
    logits = attn.experts.router.layer.output.save()   # [batch * seq, num_experts]
with model.trace(prompt):
    heads = attn.attention_head_outputs.save()         # [batch, seq, num_heads, head_dim]
    out = attn.attention_output.save()

k, kv = model.config.num_experts_per_tok, model.num_kv_heads
top = logits.topk(k, dim=-1)
gates, idx = top.values.softmax(-1), top.indices      # [batch * seq, k]
experts = attn._module.experts
W = experts.output_linear.weight[idx]                 # each slot's expert
slots = heads.reshape(-1, k, kv * model.head_dim, 1)  # each slot's 16 heads
terms = (W @ slots).squeeze(-1)                       # [batch * seq, k, hidden_size]
manual = (gates[..., None] * terms).sum(1) + experts.bias
torch.testing.assert_close(manual.view_as(out), out)

attention_keys and attention_values have num_heads heads, two copies of the 16. Each copy is its own slice of the tensor, so an edit to head h of attention_keys reaches slot h // 16's queries only; edit key/value head j at heads j and j + 16, or edit kv_proj.output, to reach every query that reads it. The expert behind a query head changes from token to token, so a query head index names a slot, not an expert.

Ablating an attention expert

Zeroing the head outputs of the slots that chose expert e removes exactly that expert's term from attention_output at those tokens, since its output projection has no bias; the other slot's term and experts.bias stay. attention_output changes on exactly the tokens whose attention router chose e:

attn, e = model.layers[1].self_attn, 3
k, kv = model.config.num_experts_per_tok, model.num_kv_heads
with model.trace(prompt):
    logits = attn.experts.router.layer.output.save()
    clean = attn.attention_output.save()
idx = logits.topk(k, dim=-1).indices.view(*clean.shape[:2], k)   # [batch, seq, k]
with model.trace(prompt):
    for s in range(k):
        attn.attention_head_outputs[:, :, s * kv:(s + 1) * kv][idx[:, :, s] == e] = 0
    ablated = attn.attention_output.save()

changed = (ablated != clean).any(-1)
assert torch.equal(changed, (idx == e).any(-1))

Loading: eager for the attention interior

The attention runs on transformers' shared interface: the six interior values need attn_implementation="eager", and the default sdpa load computes the same attention without them. The head width is kv_channels (128), so head_dim is 128 while hidden_size / num_heads is 64: the 32 query heads are 4096 wide in all, twice the stream.

The MLP: a top-2 softmax router, a bias, and a read order

The MLP's router (mlp.router) takes the top 2 of its 8 logits (mlp.router.layer.output) and a softmax over those 2, so expert_weights sum to one. It then sorts the slots by expert and the mixture runs input_linear (gate and up, 2 × intermediate_size wide, the first half through SiLU) and output_linear over the sorted slots, so expert_outputs is unavailable. expert_indices and expert_weights are read before the sort, in token order, and routed_output is the routed sum before the mixture adds its bias:

moe = model.layers[1].mlp
with model.trace(prompt):
    routed = moe.routed_output.save()
    out = moe.mlp_output.save()

torch.testing.assert_close(out - routed, moe._module.bias.expand_as(out))

The router computes the indices before the weights, so a trace reads router_logits, then expert_indices, then expert_weights. An ablation reads the indices first:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    hit = moe.expert_indices == e                              # indices first
    moe.expert_weights = moe.expert_weights.masked_fill(hit, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The readout and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. lm_head and embed_tokens share one weight. The tokenizer prepends <s> (id 1), so model.input_ids starts with it.

From the family module

04 · nnterp/families/jetmoe.py

JetMoE (JetMoeForCausalLM).

Llama's tree with the attention under its own name: model.{embed_tokens, layers[i].{input_layernorm, self_attention, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, which returns a bare tensor. self_attention is aliased self_attn.

The attention is a mixture of attention heads: a router picks num_experts_per_tok experts per token, each expert projects that token's queries for num_key_value_heads heads (experts.map) and later projects their outputs back (experts.reduce, in place of an o_proj). Keys and values come from one shared kv_proj and are tiled (repeat, not interleaved) to every query head before the shared eager forward, so the interface sees num_heads key/value heads and interface head h is routing slot h // num_kv_heads over kv head h % num_kv_heads: which expert's parameters a head's query came from varies by token. The attention returns (attn_output, attn_weights, router_logits). The MLP is a mixture of experts (Moe) that returns the routed hidden states plus a bias as one tensor.

The MLP's mixture has no experts module: its router (JetMoeTopKGating) takes a top-k of the logits of a child nn.Linear (router.layer), softmaxes the top-k, and sorts the slots by expert; the mixture then runs input_linear and output_linear over the sorted slots. So expert_indices and expert_weights are the router's top-k indices and gates, [tokens, top_k] in token order, before the sort; routed_output is the mixture's routed sum before + bias; and the per-slot outputs, sorted by expert, are unavailable.