Add norm before up projection for LatentMoE

[i-riyad]

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

What does this PR do?

Add an optional normalization layer after LatentMoE expert-output combination and before the latent up-projection:

experts → combine → fc2_norm → fc2_latent_proj → transformer hidden size

The option is exposed as moe_use_norm_before_up_proj and remains disabled by default, so existing LatentMoE models are unchanged.

When enabled, TENorm selects LayerNorm or RMSNorm from the existing transformer configuration and uses moe_latent_size as the normalized dimension. The LatentMoE projection path already requires Transformer Engine.

This PR also restores moe_use_norm_before_up_proj from checkpoint arguments alongside moe_latent_size. This ensures checkpoint loading reconstructs the norm module and consumes saved mlp.fc2_norm.weight tensors. Without restoring the option, strict checkpoint auditing reports those tensors as unconsumed because the runtime model does not construct fc2_norm.

Regression coverage exercises normalized LatentMoE execution across dispatcher, expert implementation, grouped-GEMM, and latent-size combinations, plus checkpoint-argument restoration.

Related PR: #6449 proposes similar normalized-LatentMoE behavior under the configuration name moe_latent_output_norm and targets dev. This PR retains moe_use_norm_before_up_proj because that key is persisted in existing GDP checkpoint metadata. The naming or compatibility approach can be aligned during review.

Issue tracking

Linked issue: N/A — compatibility fix for existing LatentMoE checkpoints.

Testing

• BASE_REF=main CHECK_ONLY=true SKIP_DOCS=false bash tools/autoformat.sh
  • Black, isort, pylint, and ruff passed.
• Focused distributed unit tests on four GPUs, Slurm job 557079:
  • TestLatentMoELayer::test_latent_moe_layer across all 12 parameter combinations
  • test_load_args_restores_latent_moe_norm_from_checkpoint across both parameter combinations
  • Relevant result: 14 norm/checkpoint cases passed per rank.
• The checkpoint restoration path was also exercised while loading the phase1_v2 GDP/MIMO checkpoints on the staging branch.

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

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate review.

All PRs start as draft. Mark this PR ready only after CI is passing.

[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.