[ENHANCEMENT] Optional output normalization for latent MoE
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Is your feature request related to a problem? Please describe.**
Megatron Core already supports latent MoE through `moe_latent_size` (down-projection → routed experts → combine → up-projection), but the combined latent output is fed into `fc2_latent_proj` with no way to normalize it. The Kimi K3 technical report (§2.3.1 "Normalized LatentMoE", Eq. 11) shows this step is a necessary part of the latent-MoE pipeline: the aggregated routed representation varies in scale with expert selection and routing weights, and normalizing it before the up-projection stabilizes training and improves validation loss. Today, models that need this behavior cannot express it with the generic `MoELayer`.
Tagging @NVIDIA/mcore-oncall for visibility.
**Describe the solution you'd like**
Add an optional, architecture-agnostic `moe_latent_output_norm: bool = False` field to `TransformerConfig`. When enabled, `MoELayer` applies a `TENorm` over `moe_latent_size` to the combined routed-expert output, exactly after dispatcher combine and before `fc2_latent_proj`:
```
hidden_states → fc1_latent_proj → router → experts → combine
→ routed_expert_norm (new, optional) → fc2_latent_proj → output
```
The norm type and epsilon follow the existing `normalization` / `layernorm_epsilon` configs. The feature is disabled by default, so existing latent-MoE and standard MoE models are unaffected, and it works uniformly across token dispatchers (allgather / alltoall) and expert-parallel layouts.
**Describe alternatives you've considered**
Implementing the normalization in model-specific wrapper code (e.g., a custom `MoELayer` subclass in Megatron-Bridge). However, the normalization point sits between dispatcher combine and `fc2_latent_proj` inside `MoELayer.postprocess` — it is part of the layer's execution semantics, so a wrapper must reimplement or monkey-patch `postprocess`, duplicating core logic and drifting out of sync with upstream changes. NVIDIA-NeMo/Megatron-Bridge#5130 is a concrete instance of this: its `KimiK3MoELayer(MoELayer)` overrides `postprocess` solely to insert this normalization, copying the rest of the method.
**Additional context**
Kimi K3 is the first public consumer of this option; the implementation itself contains nothing K3-specific. Landing this option in Core would let NVIDIA-NeMo/Megatron-Bridge#5130 drop its `postprocess` override. A PR implementing this feature is ready and will reference this issue.
Contributor guide
Research direction
Start with TransformerConfig and MoELayer.postprocess, especially the point after dispatcher combine and before fc2_latent_proj. Check how existing normalization and layernorm_epsilon settings are used, then verify that the optional path is disabled by default and applies consistently across the stated dispatchers and expert-parallel layouts.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100