Fix MoE aux loss tracker hang with MTP enabled

[Victarry]

Mirror of #3400

Summary

Fix a deadlock (NCCL hang) in reduce_aux_losses_tracker_across_ranks when MTP (Multi-Token Prediction) is enabled with MoE and pipeline parallelism.

Bug Description

When MTP is enabled together with MoE and pipeline parallelism (--mtp-num-layers 1), the training hangs at the first training_log call. All ranks are stuck in an all_reduce inside reduce_aux_losses_tracker_across_ranks.

Root cause: The MoE aux loss tracker values tensor has mismatched sizes across PP ranks, causing the NCCL all_reduce to hang:

• During the forward pass, the router (router.py) creates the tracker values tensor with size num_layers + mtp_num_layers (e.g. 14 + 1 = 15).
• During logging, track_moe_metrics() with force_initialize=True creates the tensor with size num_layers only (e.g. 14), but only for tracker keys that don't already exist.
• PP ranks that have no MoE layers (e.g., the first PP stage with only dense layers) never populate the tracker during forward, so force_initialize creates a size-14 tensor.
• PP ranks that have MoE layers already have size-15 tensors from the router.
• The all_reduce across the PP group then receives mismatched tensor sizes (14 vs 15), causing NCCL to hang indefinitely.

Repro: Any config with MoE + MTP + PP where at least one PP rank has no MoE layers. For example:

--num-layers 14 --moe-layer-freq "([0]*3+[1]*11)" --mtp-num-layers 1 \
--pipeline-model-parallel-size 4 \
--decoder-first-pipeline-num-layers 2 --decoder-last-pipeline-num-layers 2

Fix

In track_moe_metrics(), account for mtp_num_layers when creating the values tensor in the force_initialize path, matching the size used by the router in save_to_aux_losses_tracker().

Test plan

• Verified fix with DSv3 proxy config (PP=4, EP=2, 8x H100) - training completes successfully
• Verified no regression without MTP enabled

Made with Cursor

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

[Victarry]

/ok to test d28506e

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22591284846