🐛 CI failure: test_aux_loss.py::TestRouterAuxLoss::test_aux_loss_fusion_equivalence[seq_aux_loss]
@HaochenYuan is already working on this.
Since May 27, 2026.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
CI test `tests/unit_tests/transformer/moe/test_aux_loss.py::TestRouterAuxLoss::test_aux_loss_fusion_equivalence[seq_aux_loss]` failed in job [`tests/unit_tests/transformer/moe/**/*.py - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/26467462352/job/77944314720).
Tag the [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall) to get oncall's attention to this issue.
**Failing run**
| Field | Value |
|-------|-------|
| PR | [#4992: chore: Update transformer-engine dependency to version 2.16.0](https://github.com/NVIDIA/Megatron-LM/pull/4992) |
| Run | [26467462352](https://github.com/NVIDIA/Megatron-LM/actions/runs/26467462352) |
| Job | [tests/unit_tests/transformer/moe/**/*.py - latest](https://github.com/NVIDIA/Megatron-LM/actions/runs/26467462352/job/77944314720) |
**Error**
```text
tests/unit_tests/transformer/moe/test_aux_loss.py::TestRouterAuxLoss::test_aux_loss_fusion_equivalence[seq_aux_loss] FAILED
E RuntimeError: /root/.cache/uv/git-v0/checkouts/f679917ede501d2d/289e5a6b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu:90 in function fused_moe_aux_loss_forward_kernel_launcher: Assertion failed: num_experts == num_cols. Number of experts (32) must be equal to number of input columns (64).
../venv/lib/python3.12/site-packages/transformer_engine/pytorch/router.py:241: RuntimeError
FAILED tests/unit_tests/transformer/moe/test_aux_loss.py::TestRouterAuxLoss::test_aux_loss_fusion_equivalence[seq_aux_loss] - RuntimeError: /root/.cache/uv/git-v0/checkouts/f679917ede501d2d/289e5a6b/tr...
= 1 failed, 340 passed, 40 skipped, 200 deselected, 1395 warnings in 661.89s (0:11:01) =
```
**Steps/Code to reproduce bug**
Re-run the failing CI job linked above, or locally inside the dev container:
```bash
pytest tests/unit_tests/transformer/moe/test_aux_loss.py::TestRouterAuxLoss::test_aux_loss_fusion_equivalence[seq_aux_loss]
```
**Additional context**
Umbrella issue: #5027.
Triaged automatically via `/triage-issue`.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.