[BUG] Hybrid MTP1 with MoE fails during loss logging
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**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](https://github.com/NVIDIA/Megatron-LM/blob/c6be9199750845bd99a7d1a438862bb9ab29b630/megatron/core/transformer/moe/router.py#L596-L609) 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](https://github.com/NVIDIA/Megatron-LM/pull/4942) and [#4798](https://github.com/NVIDIA/Megatron-LM/pull/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
Contributor guide
Assessment
This issue has not been assessed yet.