[Bug] Using q instead of normalized q in Megatron's DSA MLA indexer for GLM 5 models

[SamuelGabriel]

Bug Description

In slime_plugins/models/glm5/glm5.py, the DSA lightning-indexer projects its query from the raw q_a_proj output (q_compressed), i.e. it skips the q_a_layernorm RMSNorm. The reference GLM-5 / DeepSeek-V3.2 implementations (HF transformers and vLLM) feed the indexer the q_a_layernorm-normalized q.

I've run an experiment to see the difference of doing this or not and couldn't see a significant difference in terms of abs diff between the logprobs of vLLM and megatron, but wanted to raise this still, since I believe this should actually change the top k and thus make the model a little bit off policy.

Locations in the code

I've let AI generate views of where this happens (or rather not) in slime vs HF reference to find it easily in the code:

Where slime skips the norm

q_layernorm is IdentityOp, so the q-a-layernorm is fused only into the main linear_q_up_proj; the indexer reads the pre-norm q_compressed:

# slime_plugins/models/glm5/glm5.py
q_compressed, _ = self.linear_q_down_proj(hidden_states)   # L537: raw q_a_proj output
...
q_compressed = self.q_layernorm(q_compressed)              # L551: q_layernorm = IdentityOp -> no-op
q, _ = self.linear_q_up_proj(q_compressed)                 # L552: main path normalizes internally (fused LN)
...
q_compressed = q_compressed.detach()                       # L616: still un-normalized
index_q, _ = self.wq_b(q_compressed)                       # L620: indexer query from RAW q

• slime/slime_plugins/models/glm5/glm5.py

  Lines 616 to 620 in fa3c990

  	q_compressed = q_compressed.detach()
  	hidden_states = hidden_states.detach()
  	rotary_pos_emb = rotary_pos_emb.detach()
  	index_q, _ = self.wq_b(q_compressed)

• spec sets q_layernorm=IdentityOp:

  slime/slime_plugins/models/glm5/glm5.py

  Line 764 in fa3c990

  q_layernorm=IdentityOp,

Reference (HF transformers)

GlmMoeDsaAttention computes q_resid = q_a_layernorm(q_a_proj(hidden)) and passes that to the indexer, whose wq_b consumes it:

# transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))   # L423
...
topk_indices = self.indexer(hidden_states, q_resid, ...)     # L446

# class GlmMoeDsaIndexer.forward:
#   q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`  (L220, docstring)
q = self.wq_b(q_resid)                                       # L231

• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L423
• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L446
• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L220-L231

Steps to Reproduce

Just any normal GLM 5 training.

Expected Behavior

We are very close to expected behavior, this is more of a slight imprecision.

Actual Behavior

No normalization, where HF normalizes.

Environment

none

Logs

Additional Context

No response

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.