NVIDIA / NVIDIA/Megatron-LM

🐛 Flaky CI: test_gumbel_determinism (FlextronRouter hard_sample uses un-seeded random)

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

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.