[IMPORTANT Correctness] hc_head_fn, hc_head_base, and hc_head_scale are bare nn.Parameter objects, not nn.Module children. TransformerBlock.sharded_state_dict() has custom logic that only iterates self.named_children() (layers + final_layernorm, etc.) — it does not call self._save_to_state_dict() like the base MegatronModule.sharded_state_dict() does. As a result, these three learned parameters will be silently dropped during distributed checkpoint save and won't be restored on resume, causing training divergence after checkpoint reload.

The same parameters on MultiTokenPredictionLayer are fine because its sharded_state_dict calls super().sharded_state_dict(), which does handle bare parameters.

Fix: Add explicit handling for these parameters in TransformerBlock.sharded_state_dict(). For example, after the named_children loop, call self._save_to_state_dict(...) for the standalone parameters, or wrap them in a small nn.Module so they're picked up by named_children().

@hxbai this one make sense and I think without wrapping those are not handled correctly by dist ckpting. These parameters are meant to be loaded on restart, right?

fixed

[SUGGESTION Simplification] _postprocess is called here to contract each MTP layer's multi-stream output for the loss list. This works correctly, but it means learned_output_contract + final_layernorm are computed for each MTP depth even though the multi-stream tensor was already computed in the same iteration.

Consider caching the contracted result inside _proj_and_transformer_layer (compute it but don't apply it to the returned tensor) to avoid the redundant forward through _postprocess. Not a correctness issue, but calling _postprocess externally on a "raw" tensor that the layer itself chose not to postprocess creates a subtle coupling between the caller and the layer's internal mhc_enabled flag.

[SUGGESTION Simplification] The 1e-6 offset added to pre (line 107) prevents the sigmoid gating weights from being exactly zero. However, this also means the minimum contribution per stream is always positive (1e-6), which accumulates to n * 1e-6 over all streams. This is negligible for typical hidden sizes, but it's worth noting that eps (the RMS epsilon) and this 1e-6 serve different numerical stability purposes — consider making the sigmoid floor configurable or documenting its purpose in the docstring.

Also, head_fn has shape [n, n*h] which means the F.linear projects the full concatenated multi-stream vector down to n mixing weights. This is an O(n² * h) operation per token. For large n, this could be meaningful. A comment noting this design choice (full-rank vs. per-stream factored projection) would help future readers.

replaced with eps

[SUGGESTION Simplification] This mscale = 1.0 override is duplicated in three places across two files (csa.py and deepseek_v4_hybrid_attention.py ×2). All three share the identical comment and logic pattern. Consider handling this at a single point — for example, in the rotary embedding module's constructor or a config-driven flag — rather than patching the returned value at every call site. This would also prevent future call sites from forgetting the override.

This is the special behavior of DSv4 and no need to put it into a standalone position.

[SUGGESTION Simplification] This hc_head_* parameter initialization block is duplicated verbatim in TransformerBlock.__init__ (lines 388–398). Consider extracting a small helper (e.g., _init_learned_contraction_params(config) returning a ParameterDict or small nn.Module) to avoid the duplication and ensure both sites stay in sync.