NVIDIA / NVIDIA/Megatron-LM

[BUG] Hybrid MTP1 with MoE fails during loss logging

Open
#7,172 0 comments 0 reactions 1 assignee Claimed by @Phlip79 View on GitHub
community-request waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.