Flaky: TestFusedApplyMLARope::test_forward_backward_for_q[thd] backward mismatch exceeds bf16 tolerances
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
### Summary
`tests/unit_tests/fusions/test_mla_yarn_rope_apply.py::TestFusedApplyMLARope::test_forward_backward_for_q[thd]` is intermittently failing in CI with a backward-pass numerical mismatch that exceeds bf16 tolerances.
### Observed failure
CI run: https://github.com/NVIDIA/Megatron-LM/actions/runs/25408640123 (job `tests/unit_tests/**/*.py - latest`, ID `74525633131`)
```
FAILED tests/unit_tests/fusions/test_mla_yarn_rope_apply.py::TestFusedApplyMLARope::test_forward_backward_for_q[thd]
E AssertionError: Mismatch in bwd: Tensor-likes are not close!
E Mismatched elements: 31 / 786432 (0.0%)
E Greatest absolute difference: 3.015625 at index (104, 29, 170) (up to 0.05 allowed)
E Greatest relative difference: 33.509803771972656 at index (104, 28, 166) (up to 0.02 allowed)
```
The assertion at `tests/unit_tests/fusions/test_mla_yarn_rope_apply.py:111` compares the backward gradient of the reference `apply_rotary_pos_emb` against the fused `fused_apply_mla_rope_for_q` in bf16, `thd` packed-sequence layout, `cu_seqlens=[0, 27, 54, 99, 128]`.
### Why this is non-deterministic
- Only 31 / 786,432 elements (~0.004%) exceed tolerance.
- bf16 tolerances: `atol=5e-2`, `rtol=2e-2`.
- The same merge-queue commit passed this job on rerun in workflow [25415024224](https://github.com/NVIDIA/Megatron-LM/actions/runs/25415024224).
- The fused MLA YARN RoPE backward kernel produces small numerical drift in the `thd` (packed-sequence) path that occasionally exceeds bf16 tolerances at outlier indices.
### Owning code
- Test: `tests/unit_tests/fusions/test_mla_yarn_rope_apply.py`
- Kernel: `megatron/core/fusions/fused_mla_yarn_rope_apply.py`
- Introduced in MR `!2949` (perf(mla, experimental): MLA RoPE fusion and YARN embedding cache); removed experimental tag in #2233.
### Mitigation
Marked `flaky_in_dev` in #4639 as a stop-gap. The underlying numerical drift in the fused backward kernel for `thd` should be investigated and tightened so the test can be re-enabled.
Contributor guide
Assessment
This issue has not been assessed yet.