NVIDIA / NVIDIA/Megatron-LM

🐛 CI failure: TestFusedDSAIndexerLossGradient::test_fused_indexer_loss_gradient_matches_autograd[True] — grad mismatch on 26.04

Open
#5,054 0 comments 0 reactions 1 assignee Claimed by @hxbai View on GitHub
bug
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

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.