nnterp · encyclopedia glm_moe_dsa · nnterp.families.glm_moe_dsa

family glm_moe_dsa

GLM-5

Latent attention over only the keys an indexer selects for each query, a selection each block returns beside the stream for the next block to reuse, and a sigmoid-routed mixture with a shared expert after three dense blocks.

Latent attentionSparse attentionInterleaved rotaryMixture of expertsDense first blocks

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× 78
→
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 GlmMoeDsaForCausalLMlayoutwhat 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] GlmMoeDsaDecoderLayerlayoutwhat 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 GlmMoeDsaAttentionlayoutwhat 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 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 GlmMoeDsaMLP / GlmMoeDsaMoElayoutwhat 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

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 child module shared_experts's output

Printout

What print(model) shows on zai-org/GLM-5: the native modules, standard names and every standard value with its layout.

GlmMoeDsaForCausalLM(
  (lm_head): Linear(in_features=6144, out_features=154880, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(154880, 6144, padding_idx=154820)
  (layers): ModuleList(
    (0-2): 3 x GlmMoeDsaDecoderLayer(
      (self_attn): GlmMoeDsaAttention(
        (q_a_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_a_layernorm): GlmMoeDsaRMSNorm((2048,), eps=1e-06)
        (q_b_proj): Linear(in_features=2048, out_features=16384, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=6144, out_features=576, bias=False)
        (kv_a_layernorm): GlmMoeDsaRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=28672, bias=False)
        (o_proj): Linear(in_features=16384, out_features=6144, bias=False)
        (indexer): GlmMoeDsaIndexer(
          (wq_b): Linear(in_features=2048, out_features=4096, bias=False)
          (wk): Linear(in_features=6144, out_features=128, bias=False)
          (k_norm): LayerNorm((128,), eps=1e-06, elementwise_affine=True, bias=True)
          (weights_proj): Linear(in_features=6144, out_features=32, bias=False)
        )
        (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_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
      )
      (mlp): GlmMoeDsaMLP(
        (gate_proj): Linear(in_features=6144, out_features=12288, bias=False)
        (up_proj): Linear(in_features=6144, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=6144, bias=False)
        (act_fn): SiLUActivation()
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): GlmMoeDsaRMSNorm((6144,), eps=1e-05)
      (post_attention_layernorm): GlmMoeDsaRMSNorm((6144,), eps=1e-05)
      (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
    )
    (3-77): 75 x GlmMoeDsaDecoderLayer(
      (self_attn): GlmMoeDsaAttention(
        (q_a_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_a_layernorm): GlmMoeDsaRMSNorm((2048,), eps=1e-06)
        (q_b_proj): Linear(in_features=2048, out_features=16384, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=6144, out_features=576, bias=False)
        (kv_a_layernorm): GlmMoeDsaRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=28672, bias=False)
        (o_proj): Linear(in_features=16384, out_features=6144, bias=False)
        (indexer): GlmMoeDsaIndexer(
          (wq_b): Linear(in_features=2048, out_features=4096, bias=False)
          (wk): Linear(in_features=6144, out_features=128, bias=False)
          (k_norm): LayerNorm((128,), eps=1e-06, elementwise_affine=True, bias=True)
          (weights_proj): Linear(in_features=6144, out_features=32, bias=False)
        )
        (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_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
      )
      (mlp): GlmMoeDsaMoE(
        (experts): GlmMoeDsaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): GlmMoeDsaTopkRouter()
        (shared_experts): GlmMoeDsaMLP(
          (gate_proj): Linear(in_features=6144, out_features=2048, bias=False)
          (up_proj): Linear(in_features=6144, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=6144, bias=False)
          (act_fn): SiLUActivation()
          (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
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
        )
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): GlmMoeDsaRMSNorm((6144,), eps=1e-05)
      (post_attention_layernorm): GlmMoeDsaRMSNorm((6144,), eps=1e-05)
      (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): GlmMoeDsaRMSNorm((6144,), 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_layers78
model.hidden_size6144
model.num_heads64
model.num_kv_heads64
model.head_dim256
model.qk_head_dim256
model.vocab_size154880
model.intermediate_size12288
architecturesGlmMoeDsaForCausalLM
num_hidden_layers78
hidden_size6144
intermediate_size12288
num_attention_heads64
num_key_value_heads64
head_dim64
vocab_size154880
max_position_embeddings202752
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
q_lora_rank2048
kv_lora_rank512
qk_nope_head_dim192
qk_rope_head_dim64
v_head_dim256
first_k_dense_replace3
moe_intermediate_size2048
n_routed_experts256
num_experts_per_tok8
n_shared_experts1
routed_scaling_factor2.5
norm_topk_probTrue
attention_biasFalse
n_group1
topk_group1

The Notes

03 · what to know before you run

The block, in order

a, sel = self_attn(input_layernorm(x), sel)   # sparse latent attention; sel is the key selection
h      = x + a
out    = h + mlp(post_attention_layernorm(h))  # dense on blocks 0-2, a mixture after
return out, sel

Llama's pre-norm block and Llama's names. The block returns (hidden_states, topk_indices), and the attention (attn_output, attn_weights, topk_indices): the key selection travels to the next block beside the stream. layer_output and attention_output are the first elements of these tuples, and the identity is the plain sum on a dense block and on a mixture block:

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)

config.mlp_layer_types says which blocks are dense: blocks 0 to 2 on every GLM-5 checkpoint, a mixture on the other 75. The configs carry one multi-token-prediction block (num_nextn_predict_layers 1), stored after the last block; transformers does not build it, so model.layers has 78 blocks.

Latent attention: every head's key comes from one latent

attention_queries and attention_keys are qk_head_dim wide (qk_nope_head_dim + qk_rope_head_dim, 192 + 64 = 256), attention_values and attention_head_outputs head_dim (v_head_dim, also 256). transformers sets the config's head_dim to qk_rope_head_dim, the rotary width (64), whatever the checkpoint's config.json says; no served value is that wide. Queries go through a compressed rank (q_lora_rank 2048): q_a_proj, q_a_layernorm, then q_b_proj, with no q_proj module. Keys and values are expanded for every head from one latent per token (kv_lora_rank 512), so attention_keys has num_heads (64) heads. The last 64 dimensions of every head's key are one rotary key that all heads share. The scores are scaled by qk_head_dim ** -0.5 (self_attn.scaling, 0.0625).

attn = model.layers[1].self_attn
with model.trace(prompt):
    q = attn.attention_queries.save()          # [1, 64, seq, 256]
    k = attn.attention_keys.save()             # [1, 64, seq, 256]
    v = attn.attention_values.save()           # [1, 64, seq, 256]

assert torch.equal(k[:, 0, :, -64:], k[:, 1, :, -64:])   # one rotary key

The rotary part is served de-interleaved

The attention rotates pair (2i, 2i + 1) of the projection into dimensions i and i + 32 of the rotary part on every checkpoint. The forward always does this, whatever rope_interleave says. attention_queries[..., -64:] and attention_keys[..., -64:] are in that rotate-half order: at position 0, where rotary is the identity, the served rotary part is q_b_proj's even dimensions followed by its odd ones.

with model.trace(prompt):
    proj = attn.q_b_proj.output.save()
with model.trace(prompt):
    q = attn.attention_queries.save()

rope = proj.view(1, -1, model.num_heads, model.qk_head_dim)[0, 0, 0, -64:]
torch.testing.assert_close(q[0, 0, 0, -64:], torch.cat([rope[0::2], rope[1::2]]))

Up to 2048 tokens the sparse attention is plain causal attention

A lightning indexer (self_attn.indexer: index_n_heads 32 heads of index_head_dim 128) scores every key for every query and keeps the index_topk (2048) best. The attention sets the other keys' mask entries to the dtype's minimum before it calls the attention function. It does this under eager and under the default sdpa load alike, so both loads run the same sparse computation; the interior values need attn_implementation="eager". attention_scores carry that mask, so a dropped key reads like a future one. attention_probabilities is the dense [batch, num_heads, query, key] pattern, exactly zero outside each query's selection.

The selection is self_attn.output[2] on every block, int32 key indices of shape [batch, query, min(index_topk, seq)]; it has no standard name. On a prompt of up to 2048 tokens it lists every key, the causal mask removes the future ones, and the pattern is plain causal: row t has t + 1 nonzero entries.

with model.trace(prompt):
    probs = attn.attention_probabilities.save()
    sel = attn.output[2].save()                # [1, seq, min(2048, seq)]

rows = (probs[0, 0] != 0).sum(-1)
assert torch.equal(rows, torch.arange(1, rows.numel() + 1))   # plain causal

A written pattern is used as written, so a write can give weight to a key the indexer dropped. The indexer runs under torch.no_grad(): no gradient reaches its projections, and the selection does not move with a gradient-based attribution.

The selection is shared between blocks on GLM-5.2 and GLM-5.3

config.indexer_types marks each block "full" (it has an indexer and selects its own keys) or "shared" (no indexer; it reuses the selection the block before it returned). Every block is full on GLM-5 and GLM-5.1. On GLM-5.2 and GLM-5.3, 21 blocks are full (0, 1, 2, then every fourth from 6) and the 57 others shared: blocks 7 to 9 attend with block 6's selection. skip_layers hands the block after a skipped range no selection, which a shared block refuses with a ValueError: skip up to a full block or to the end.

The mixture: sigmoid scores, a selection bias, renormalized weights

256 routed experts, 8 per token, and one shared expert every token runs through. router_logits are the logits before the sigmoid, and before the selection bias (router.e_score_correction_bias). The bias is added to the sigmoids only to choose; n_group and topk_group are 1, so every expert is a candidate. expert_weights are the chosen experts' sigmoids without the bias, renormalized to sum to one (norm_topk_prob) and multiplied by routed_scaling_factor, so they sum to 2.5 at every token.

moe = model.layers[3].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()
    w = moe.expert_weights.save()
    idx = moe.expert_indices.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    out = moe.mlp_output.save()

chosen = logits.sigmoid().gather(-1, idx)
scale = model.config.routed_scaling_factor
torch.testing.assert_close(w, chosen / chosen.sum(-1, keepdim=True) * scale)
torch.testing.assert_close(routed + shared, out)

The shared expert runs after the routed ones, so shared_expert_output is read after routed_output. Ablating an expert is zeroing the weight of every slot that chose it; the other slots keep their weights, so the token's weights then sum to less than 2.5:

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

Blocks 0 to 2 have a dense MLP: support() reports every mixture value missing there.

From the family module

04 · nnterp/families/glm_moe_dsa.py

GLM-5 (GlmMoeDsaForCausalLM).

DeepSeek-V3.2's tree and block (model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block): multi-head latent attention with DeepSeek Sparse Attention, the indexer's top-k selection folded into the attention mask before the shared attention_interface call under eager and sdpa, so the interior is the base's and means what it means on deepseek_v32: attention_scores carry the sparse mask, attention_probabilities is the dense [batch, num_heads, query, key] pattern, zero outside each query's selected keys. While a prompt is no longer than index_topk the selection is every causal key.

What differs is that the selection is shared across blocks. config.indexer_types marks each block "full" (it has an indexer and selects its own keys) or "shared" (no indexer; it reuses the selection of the block before). The selection therefore travels between blocks: the attention returns (attn_output, attn_weights, topk_indices) and the block returns (hidden_states, topk_indices), so returns_tuple is set; attention_output and layer_output take the first element. The selection is self_attn.output[2] on every block (self_attn.indexer.output on a full one); it has no standard name. skip_layers hands a skipped block's successor None for the selection, which a shared block refuses (transformers raises ValueError), so skip up to a full block or to the end. Dense and mixture-of-experts MLPs follow mlp_layer_types. The sizes are DeepSeek-V2's.