[BUG] Hybrid MTP1 with MoE fails during loss logging

Author: jeyaletheiaCreated Sep 9, 2026Updated Sep 18, 2026
Labelscommunity-request

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:

python
hybrid_layer_pattern = "*E*E/*E"
mtp_num_layers = 1
moe_aux_loss_coeff = 0.01
moe_z_loss_coeff = 0.001

The 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