NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Optional output normalization for latent MoE

Open
#6,448 0 comments 0 reactions 0 assignees View on GitHub
community-request enhancement waiting-on-maintainers
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.