[MoE] Add optional output normalization for latent MoE

[86MaxCao]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add an optional normalization layer on the combined routed-expert output in latent MoE, applied after token combination and before the up-projection back to the transformer hidden size.

Motivation

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. Since the normalization point sits between dispatcher combine and fc2_latent_proj inside MoELayer's execution flow, it belongs in the layer itself rather than in model-specific wrapper code.

This PR exposes it as an optional, architecture-agnostic configuration that is disabled by default. Kimi K3 is the first public consumer; nothing in the implementation is K3-specific, and it works uniformly across token dispatchers (allgather / alltoall) and expert-parallel layouts.

A concurrent Megatron-Bridge integration (NVIDIA-NeMo/Megatron-Bridge#5130) currently duplicates MoELayer.postprocess to insert exactly this normalization; landing the option in Core would let that integration drop the override.

Changes

New config field moe_latent_output_norm: bool = False. When enabled, the combined routed-expert output is normalized in the latent dimension after combine and before the up-projection:

hidden_states → fc1_latent_proj → router → experts → combine
→ routed_expert_norm (new, optional) → fc2_latent_proj → output

• transformer_config.py: add moe_latent_output_norm field + validation requiring moe_latent_size.
• moe_layer.py: build routed_expert_norm = TENorm(...) when enabled; apply it in postprocess between combine_postprocess and fc2_latent_proj. The norm type and epsilon follow the existing normalization / layernorm_epsilon configs.
• arguments.py / checkpointing.py: argument validation, pass-through, and checkpoint restore. The CLI flag --moe-latent-output-norm is auto-generated and requires --moe-latent-size.
• test_latent_moe_layer.py: enable the option in the layer test; add config-validation and execution-order tests; use latent sizes 64/128.

Test plan

• pytest tests/unit_tests/transformer/moe/test_latent_moe_layer.py (14 passed: allgather/alltoall × TE/local × grouped/non-grouped GEMM × latent size 64/128, plus config-validation and execution-order tests)
• End-to-end training convergence check with --moe-latent-output-norm

Issue tracking

Linked issue: Fix #6448

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.