[BUG] Hybrid MTP1 with MoE fails during loss logging
Describe the bug
On upstream c6be9199750845bd99a7d1a438862bb9ab29b630, HybridModel with MTP1 and an inner MoE raises IndexError when router auxiliary/z losses are enabled. Engram is not involved.
Steps/Code to reproduce bug
Run forward/backward with hidden size 32, four attention heads, four experts/top-2, zero dropout, and:
hybrid_layer_pattern = "*E*E/*E"
mtp_num_layers = 1
moe_aux_loss_coeff = 0.01
moe_z_loss_coeff = 0.001The router's metric indexing treats the inner MoE position, 2, as MTP prediction depth, which is 1. Four main Hybrid nodes plus MTP1 allocate five slots, but the computed index is 4 + 2 - 1 = 5.
| Control | Upstream | MTP fix only |
|---|---|---|
| Hybrid, router losses enabled | IndexError |
Pass |
| GPT, router losses enabled | Pass | Pass |
| Either model, router losses disabled | Pass | Pass |
Expected behavior
Training should complete normally. Record metrics by the enclosing MTP prediction depth without changing internal router/replay identities.
Additional context
I prepared a fix that passes prediction depth explicitly, with real forward/backward regression tests covering the controls above, MTP2 and multiple inner MoE nodes. The tests will accompany the fix PR.
#4942 and #4798 propose modulo-based fixes for this issue; I can coordinate and contribute the regression coverage.
Local environment: 8×H800, Python 3.12.3, PyTorch 2.12.0a0+5aff3928d8.nv26.05, Transformer Engine 2.15.0+42b84005, CUDA 13.2. These are fixed-commit local results, not official CI.
@NVIDIA/mcore-oncall
Source: NVIDIA/Megatron-LM