nnterp · encyclopedia afmoe · nnterp.families.afmoe

family afmoe

Trinity NanoTrinity MiniTrinity Large

A sandwich block whose attention gates its heads with a sigmoid and applies rotary only on its sliding-window blocks, and whose MLP is dense on the first blocks and a mixture with a shared expert after them, its routed experts picked by sigmoid scores plus a selection bias.

Sandwich normsEmbedding multiplierQuery/key normsOutput gateSliding windowBlocks without rotaryMixture of expertsDense first blocksGrouped-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, each tick's color gives its block's shape, and the diagram redraws for the block picked.

embed_tokenstoken_embeddings
→
layers× 56
→
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 AfmoeForCausalLMlayoutwhat 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] AfmoeDecoderLayerlayoutwhat 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 pre_mlp_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module pre_mlp_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 AfmoeAttentionlayoutwhat 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: the post-attention norm's output

the sibling module post_attention_layernorm's 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 AfmoeMLP / AfmoeSparseMoeBlocklayoutwhat 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 post-MLP norm's output

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

the child module router.gate's output
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 child module shared_experts's output

Printout

What print(model) shows on arcee-ai/Trinity-Nano-Base: the native modules, standard names and every standard value with its layout.

AfmoeForCausalLM(
  (lm_head): Linear(in_features=1024, out_features=200192, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(200192, 1024)
  (layers): ModuleList(
    (0-1): 2 x AfmoeDecoderLayer(
      (self_attn): AfmoeAttention(
        (q_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (k_proj): Linear(in_features=1024, out_features=256, bias=False)
        (v_proj): Linear(in_features=1024, out_features=256, bias=False)
        (o_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (q_norm): AfmoeRMSNorm((128,), eps=1e-05)
        (k_norm): AfmoeRMSNorm((128,), eps=1e-05)
        (gate_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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): AfmoeRMSNorm((1024,), eps=1e-05)
      (post_attention_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (pre_mlp_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (post_mlp_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (mlp): AfmoeMLP(
        (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
        (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
        (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-MLP norm'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
      )
      (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
      (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
      (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
    )
    (2-55): 54 x AfmoeDecoderLayer(
      (self_attn): AfmoeAttention(
        (q_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (k_proj): Linear(in_features=1024, out_features=256, bias=False)
        (v_proj): Linear(in_features=1024, out_features=256, bias=False)
        (o_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (q_norm): AfmoeRMSNorm((128,), eps=1e-05)
        (k_norm): AfmoeRMSNorm((128,), eps=1e-05)
        (gate_proj): Linear(in_features=1024, out_features=1024, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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): AfmoeRMSNorm((1024,), eps=1e-05)
      (post_attention_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (pre_mlp_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (post_mlp_layernorm): AfmoeRMSNorm((1024,), eps=1e-05)
      (mlp): AfmoeSparseMoeBlock(
        (router): AfmoeTokenChoiceRouter(
          (gate): Linear(in_features=1024, out_features=128, bias=False)
        )
        (shared_experts): AfmoeMLP(
          (gate_proj): Linear(in_features=1024, out_features=256, bias=False)
          (up_proj): Linear(in_features=1024, out_features=256, bias=False)
          (down_proj): Linear(in_features=256, out_features=1024, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-MLP norm'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
        )
        (experts): AfmoeExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-MLP norm's output
      )
      (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
      (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
      (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): AfmoeRMSNorm((1024,), 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_layers56
model.hidden_size1024
model.num_heads8
model.num_kv_heads2
model.head_dim128
model.qk_head_dim128
model.vocab_size200192
model.intermediate_size3072
architecturesAfmoeForCausalLM
num_hidden_layers56
hidden_size1024
intermediate_size3072
num_attention_heads8
num_key_value_heads2
head_dim128
vocab_size200192
max_position_embeddings131072
sliding_window2048
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
moe_intermediate_size256
num_experts128
num_experts_per_tok8
num_shared_experts1
attention_biasFalse
n_group1
topk_group1
route_scale2.826
mup_enabledTrue
num_dense_layers2

The Notes

03 · what to know before you run

The block, in order

h   = x + post_attention_layernorm(self_attn(input_layernorm(x)))
out = h + post_mlp_layernorm(mlp(pre_mlp_layernorm(h)))   # dense, then a mixture

Four RMSNorms per block, two per sublayer. post_attention_layernorm follows the attention, where Llama's same name is the norm before the MLP; the MLP's input norm is pre_mlp_layernorm. num_dense_layers blocks come first with a dense MLP (2 on Nano and Mini, 6 on Large), and every block after them has a mixture. The identity is the plain sum on both kinds:

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 contributions are the post-norms' outputs

attention_output is post_attention_layernorm's output and mlp_output is post_mlp_layernorm's; self_attn.output[0] and mlp.output are the tensors entering them. RMSNorm is scale-invariant up to eps, so scaling mlp.output leaves the stream as it was. Scale or steer the contribution itself.

with model.trace(prompt):
    model.layers[1].mlp.mlp_output[:] *= 0.5     # halves what the block adds

Inside the mixture, routed_output + shared_expert_output is mlp.output, the tensor entering post_mlp_layernorm, so an edit to one expert reaches the stream only after the norm has rescaled the token's whole MLP term.

Attention: per-head q/k norms and a sigmoid gate on every head

q_norm and k_norm are RMSNorms over each head's 128 dimensions, applied before the rotary. attention_queries and attention_keys are read after the norms and, on a sliding-window block, after the rotary. The attention also projects a gate from its input (gate_proj, one value per channel of every head) and multiplies the heads' outputs by the gate's sigmoid before o_proj. attention_head_outputs are the heads' outputs before the gate:

attn = model.layers[1].self_attn
with model.trace(prompt):
    gate = attn.gate_proj.output.save()
with model.trace(prompt):
    heads = attn.attention_head_outputs.save()    # [batch, seq, num_heads, head_dim], ungated
    o_in = attn.o_proj.input.save()

torch.testing.assert_close(o_in, heads.flatten(-2) * gate.sigmoid())

The heads are grouped: 8 query heads over 2 key/value heads on Nano, 32 over 4 on Mini and 48 over 8 on Large. The attention interior needs attn_implementation="eager"; the default load is sdpa.

Rotary runs on the sliding-window blocks only

Every fourth block attends over the whole prefix (global_attn_every_n_layers 4: blocks 3, 7, 11, ...); the others attend over the last sliding_window tokens (2048 on Nano and Mini, 4096 on Large). Only the sliding blocks apply the rotary to queries and keys. On a full block only the causal mask orders the tokens, and attention_queries is q_norm's output unchanged:

attn = model.layers[1].self_attn                  # a full block on the pinned tiny
with model.trace(prompt):
    qn = attn.q_norm.output.save()                # [batch, seq, num_heads, head_dim]
with model.trace(prompt):
    q = attn.attention_queries.save()             # [batch, num_heads, seq, head_dim]

torch.testing.assert_close(q, qn.transpose(1, 2))

The router: sigmoid scores, a selection bias, weights summing to route_scale

128 routed experts, 8 per token, on Nano and Mini; 256, 4 per token, on Large. router_logits are router.gate's output, in the model's dtype: the router casts them to float32 after this point, then takes the sigmoid and adds the bias. The bias, mlp.expert_bias, is a frozen parameter of the mixture; it is added to the sigmoids only to choose the experts, and there are no expert groups. expert_weights are the chosen experts' sigmoids without the bias, renormalized to sum to one and multiplied by route_scale, so they sum to 2.826 at every token on Nano and Mini and 2.448 on Large:

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()
    w = moe.expert_weights.save()
    idx = moe.expert_indices.save()

chosen = logits.float().sigmoid().gather(-1, idx)
scale = model.config.route_scale
torch.testing.assert_close(w, chosen / chosen.sum(-1, keepdim=True) * scale)

A written router_logits moves the choice and the weights; a written expert_bias moves only the choice.

The shared expert runs between the router and the routed experts

One shared expert (mlp.shared_experts, an AfmoeMLP of width moe_intermediate_size × num_shared_experts) runs on every token after the router and before the routed experts, and the mixture returns their sum. In one trace read router_logits, then shared_expert_output, then expert_weights and expert_indices (they are the routed experts' arguments), then expert_outputs and routed_output:

with model.trace(prompt):
    logits = moe.router_logits.save()
    shared = moe.shared_expert_output.save()
    w = moe.expert_weights.save()
    routed = moe.routed_output.save()
    total = moe.output.save()

torch.testing.assert_close(routed + shared, total)

nnterp serves the shared expert as an MLP too, but its mlp_output is unavailable: what the block adds is the mixture's post-normed sum, read at layers[i].mlp.mlp_output.

Ablating an expert

Zero the weight of every routing slot that chose the expert. The token's other weights keep their values, so its weights then sum to less than route_scale:

with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == 3, 0)
    ablated = model.logits.save()

routed_output changes on exactly the tokens that chose the expert. The dense first blocks have no router: model.support() reports every mixture value unavailable on them.

From the family module

04 · nnterp/families/afmoe.py

AFMoE (AfmoeForCausalLM, Arcee's Trinity).

Llama's tree with a sandwich block, every sublayer normed before and after::

h = x + post_attention_layernorm(self_attn(input_layernorm(x)))
out = h + post_mlp_layernorm(mlp(pre_mlp_layernorm(h)))

What the block adds is each post-norm's output, so attention_output points at post_attention_layernorm (which follows the attention, as on Gemma-2) and mlp_output at post_mlp_layernorm; pre_mlp_layernorm keeps its native name and mlp.input is its output. The attention runs the shared eager forward with q/k norms, the rotary only on its sliding-window (local) blocks, and a sigmoid gate (gate_proj) on the interface's output before o_proj, so attention_head_outputs are ungated. The first num_dense_layers blocks have a dense AfmoeMLP, the rest a mixture (AfmoeSparseMoeBlock) returning its shared expert's output plus the routed experts' sum as one tensor. The shared expert (mlp.shared_experts) is an AfmoeMLP too, so it is an Mlp; its mlp_output is unavailable, since the block adds the mixture's post-normed sum.