🐛 Flaky CI: test_gumbel_determinism (FlextronRouter hard_sample uses un-seeded random)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
CI test `tests/unit_tests/elastification/test_hybrid_flex_router.py::TestFlextronRouter::test_gumbel_determinism` is **flaky**. It passes on most ranks/runs but intermittently fails with:
```
AssertionError: Tensor-likes are not equal!
Mismatched elements: 2 / 2 (100.0%)
Greatest absolute difference: 0.4191093146800995 at index (1,)
Greatest relative difference: inf at index (1,)
tests/unit_tests/elastification/test_hybrid_flex_router.py:180: AssertionError
```
First observed in job [`tests/unit_tests/**/*.py - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/26937673156/job/79486685125) (rank 3 failed while rank 0 passed all 1460 tests). The triggering PR (#5152) is a CI-notification config change and is unrelated to the failure.
Tag @NVIDIA/mcore-oncall to get oncall's attention to this issue.
**Failing run**
| Field | Value |
|-------|-------|
| PR | [#5152: ci(notify): mention mcore-oncall and philipp on critical CI events](https://github.com/NVIDIA/Megatron-LM/pull/5152) |
| Run | [26937673156](https://github.com/NVIDIA/Megatron-LM/actions/runs/26937673156) |
| Job | [tests/unit_tests/**/*.py - latest](https://github.com/NVIDIA/Megatron-LM/actions/runs/26937673156/job/79486685125) |
**Root cause**
The test asserts two routers produce **bit-identical** Gumbel-softmax logits (`torch.testing.assert_close(..., atol=0, rtol=0)`). The Gumbel sample itself is deterministic — `_dp_gumbel_softmax` wraps `F.gumbel_softmax` in `torch.manual_seed(seed)` with save/restore of the RNG state.
However, `FlextronRouter.forward` derives a separate stochastic flag from Python's **global, un-seeded** RNG:
```python
# megatron/elastification/router/hybrid_flex_router.py:545
hard_sample = random.random() > self.hard_sample_th # hard_sample_th = 0.996
```
This `random.random()` draw is **not** covered by the torch seeding, so `router_a(1.0)` and `router_b(1.0)` each draw independently:
- ~99.6% of the time both get `hard_sample=False` (soft Gumbel) → identical logits → pass.
- ~0.4% of the time one router draws `> 0.996` (`hard=True`, one-hot) while the other stays soft → divergent logits → the `atol=0, rtol=0` assert fails. The `inf` relative difference is the zero entry of the one-hot vector.
It surfaces on an arbitrary rank because each rank's global Python-RNG state differs by the time this test runs.
**Proposed fix**
Make the test deterministic by seeding Python's `random` identically before each router forward, so both routers observe the same `hard_sample`:
```python
random.seed(0)
out_a = router_a(1.0)
random.seed(0)
out_b = router_b(1.0)
```
This preserves the test's intent (verifying Gumbel **RNG** determinism) without touching production behavior. A separate, deeper question for the elastification owner: whether `hard_sample` should be synchronized across DP ranks during real training (currently each process draws independently from `random.random()`).
**Mitigation (already applied)**
The test has been quarantined with `@pytest.mark.flaky_in_dev` to keep the default CI pipeline green until the fix above lands.
**Steps/Code to reproduce bug**
```bash
uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -q \
tests/unit_tests/elastification/test_hybrid_flex_router.py::TestFlextronRouter::test_gumbel_determinism
```
Repeat across several runs / ranks to observe the intermittent failure.
**Additional context**
Triaged automatically via `/create-issue`.
Contributor guide
Assessment
This issue has not been assessed yet.