🐛 CI failure: TestFusedDSAIndexerLossGradient::test_fused_indexer_loss_gradient_matches_autograd[True] — grad mismatch on 26.04
- 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/experimental_attention_variant/test_attention_variant_dsa.py::TestFusedDSAIndexerLossGradient::test_fused_indexer_loss_gradient_matches_autograd[True]` failed in job [`tests/unit_tests/transformer/**/*.py - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/26627543568/job/78473356084) under the 26.04 base-image bump on `dev`.
The fused DSA indexer-loss gradient no longer matches the autograd reference within the configured tolerance (`max diff = 8.59e-05`). Reads as numerical drift from new cuBLAS / cuDNN / TE shipped in `nvcr.io/nvidia/pytorch:26.04-py3`, not a code regression on `dev`.
Tracked under umbrella issue #5053. Tag @NVIDIA/mcore-oncall.
**Failing run**
| Field | Value |
|-------|-------|
| PR | [#5051: chore: Update Docker image version to 26.04-py3 on dev](https://github.com/NVIDIA/Megatron-LM/pull/5051) |
| Run | [26627543568](https://github.com/NVIDIA/Megatron-LM/actions/runs/26627543568) |
| Job | [`tests/unit_tests/transformer/**/*.py - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/26627543568/job/78473356084) |
**Error**
```
tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py:450: AssertionError
E AssertionError: [seqlen=32, topk=16, sparse=True, per_token=True] grad_q mismatch:
max diff = 8.594989776611328e-05
FAILED tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py::TestFusedDSAIndexerLossGradient::test_fused_indexer_loss_gradient_matches_autograd[True]
========= 1 failed, 101 passed, 20 deselected, 312 warnings in 44.20s ==========
```
**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/experimental_attention_variant/test_attention_variant_dsa.py::TestFusedDSAIndexerLossGradient::test_fused_indexer_loss_gradient_matches_autograd
```
**Additional context**
Triaged automatically via the `create-issue` skill. Likely needs a small bump to the tolerance threshold for the fused-vs-autograd grad comparison on the 26.04 container, or a re-derivation of the expected gradient reference.
Contributor guide
Assessment
This issue has not been assessed yet.