From 6b2be2c30bdc88f611166dcef2bd0cefc01edf36 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:27:32 +0000 Subject: [PATCH 1/3] Share one Megatron worker pool across the gpu_megatron modules Each of the 15 gpu_megatron modules that use `dist_workers` started its own worker pool, about 19 of the lane's 36 minutes. Fourteen logged module starts took 75-77 s each (the GDN module 179 s: 76 s pool plus 103 s Triton JIT), plus 3.6-10 s of shutdown, and the start has a p90 of 94-98 s. The 76 s is not distributed init (an in-process init takes 0.1 s): the spawned child unpickles megatron_worker_teardown by reference, which imports the heavy tests/gpu_megatron/conftest.py (transformers, megatron.core, all of modelopt and its plugins, TE) before it reaches init_process_group. The same import chain takes pytest itself 89 s to start here against 17 s in the gpu lane, whose session-scoped pool, with a light teardown function, starts in 1.7 s. Module scope was chosen when a spawn cost about 10 s ("avoids cross-test state contamination"); the git history shows no contamination incident, and the cost has grown 7.6x. _pool_cache is now session-scoped, so the pools are keyed by world size and shared by every module. The dist_workers fixtures stay function-scoped and get their pool from get_live_pool, which replaces a pool that a hang killed. What isolation the module scope gave is kept explicitly: - megatron_worker_teardown, after every job: synchronizes CUDA first (an asynchronous CUDA error then surfaces in the test that caused it), destroys Megatron's model-parallel state as before, and now also the process groups it leaves behind (about 20 NCCL groups per dispatch, each with a communicator and a watchdog thread; test_megatron.py alone dispatches about 110 times). A pool that lives through the whole session needs it more than a per-module pool did; - _reset_shared_pools, an autouse module-scoped fixture: when a module starts, it resets the process-global torch settings that helpers like set_seed change (cudnn and TF32 flags, deterministic algorithms, RNG) to what a fresh worker starts with (reset_worker_state, which the workers prepare by recording those settings right after they start), in the pools that already exist. It is not part of _get_pool, which runs once per test and would add a job to each; - get_live_pool and the pool's abort handling: a job that times out, is interrupted or loses a worker kills the pool and the next test gets a fresh one, for the cost of one pool start. That also covers CUDA poisoning after the known Blackwell illegal-memory-access flake. Expected saving: about 17 min of the gpu_megatron job (36 min of tests to about 19), and one pool start instead of 15 means fewer chances to hit the ready timeout. Testing: reset_worker_state is covered by a gloo test in tests/gpu/torch/utils/test_distributed_worker_pool.py (it runs on CPU): flags and seed changed in a job are back at their spawn-time values afterwards. The Megatron lane needs GPUs, so this PR's gpu_megatron job is the check. In an earlier CI run of the whole series (with the Triton cache of another change), the lane's pytest step took 862 s on cold caches and 336 s on warm ones, against 2,022 s for a baseline run. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 32 ++++++++++++++ .../utils/test_distributed_worker_pool.py | 37 +++++++++++++++- tests/gpu_megatron/conftest.py | 43 ++++++++++++++----- 3 files changed, 100 insertions(+), 12 deletions(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index d5bbebb5404..12a6b56595b 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -15,6 +15,7 @@ import contextlib import faulthandler +import gc import os import queue import signal @@ -96,6 +97,36 @@ def get_live_pool(pools, key, make_pool): return pool +_BASELINE_TORCH_FLAGS = {} + + +def _capture_baseline_torch_flags(): + """Remember the torch settings a freshly spawned worker starts with.""" + _BASELINE_TORCH_FLAGS.update( + cudnn_deterministic=torch.backends.cudnn.deterministic, + cudnn_benchmark=torch.backends.cudnn.benchmark, + cudnn_tf32=torch.backends.cudnn.allow_tf32, + matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + deterministic_algorithms=torch.are_deterministic_algorithms_enabled(), + ) + + +def reset_worker_state(rank, world_size): + """Return a persistent worker to the state it had right after spawning. + + ``set_seed`` and similar helpers change process-global torch settings; run this job at the + boundary between test modules that share workers. + """ + torch.backends.cudnn.deterministic = _BASELINE_TORCH_FLAGS["cudnn_deterministic"] + torch.backends.cudnn.benchmark = _BASELINE_TORCH_FLAGS["cudnn_benchmark"] + torch.backends.cudnn.allow_tf32 = _BASELINE_TORCH_FLAGS["cudnn_tf32"] + torch.backends.cuda.matmul.allow_tf32 = _BASELINE_TORCH_FLAGS["matmul_tf32"] + torch.use_deterministic_algorithms(_BASELINE_TORCH_FLAGS["deterministic_algorithms"]) + torch.manual_seed(1234) + gc.collect() + torch.cuda.empty_cache() + + def default_worker_teardown(rank, world_size): """Minimal cleanup between tests in persistent workers.""" try: @@ -207,6 +238,7 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard if backend == "nccl" and torch.cuda.is_available(): torch.cuda.set_device(rank) torch.manual_seed(1234) + _capture_baseline_torch_flags() result_queue.put("ready") while True: diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py index 3d59051c74a..e60dbc54d66 100644 --- a/tests/gpu/torch/utils/test_distributed_worker_pool.py +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -21,8 +21,9 @@ import time import pytest +import torch import torch.distributed as dist -from _test_utils.torch.distributed.utils import DistributedWorkerPool +from _test_utils.torch.distributed.utils import DistributedWorkerPool, reset_worker_state def _hang(rank, size): @@ -56,6 +57,32 @@ def _rank_0_dies_and_rank_1_waits_for_it(rank, size): time.sleep(600) +def _torch_settings(): + return ( + torch.backends.cudnn.deterministic, + torch.backends.cudnn.benchmark, + torch.backends.cudnn.allow_tf32, + torch.backends.cuda.matmul.allow_tf32, + torch.are_deterministic_algorithms_enabled(), + ) + + +def _change_the_settings_then_reset_them(rank, size): + spawn_time = _torch_settings() + torch.backends.cudnn.deterministic = not spawn_time[0] + torch.backends.cudnn.benchmark = not spawn_time[1] + torch.backends.cudnn.allow_tf32 = not spawn_time[2] + torch.backends.cuda.matmul.allow_tf32 = not spawn_time[3] + torch.use_deterministic_algorithms(not spawn_time[4]) + torch.manual_seed(5) + assert _torch_settings() != spawn_time + + reset_worker_state(rank, size) + + assert _torch_settings() == spawn_time + assert torch.initial_seed() == 1234 + + def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) try: @@ -136,3 +163,11 @@ def test_a_failed_start_raises_and_leaves_no_workers(kwargs, error): DistributedWorkerPool(world_size=1, **{"backend": "gloo", "teardown_fn": None, **kwargs}) time.sleep(1) assert {p.pid for p in multiprocessing.active_children()} <= before + + +def test_reset_worker_state_restores_the_settings_a_worker_starts_with(): + pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) + try: + pool.run(_change_the_settings_then_reset_them) + finally: + pool.shutdown() diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 51261e30023..c53a24abd49 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,7 +17,11 @@ import pytest import torch from _test_utils.fs_utils import assert_unmodified_tree -from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool +from _test_utils.torch.distributed.utils import ( + DistributedWorkerPool, + get_live_pool, + reset_worker_state, +) from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel @@ -68,6 +72,8 @@ def _prestart_preprocess_forkserver(): def megatron_worker_teardown(rank, world_size): """Clean up model-parallel state between tests in persistent workers.""" + # Surface asynchronous CUDA errors at the test that caused them, not at a later one + torch.cuda.synchronize() if dist.is_initialized(): dist.barrier() try: @@ -79,6 +85,12 @@ def megatron_worker_teardown(rank, world_size): apex_destroy() except Exception as e: print(f"Error destroying model parallel with Apex: {e}") + # destroy_model_parallel() only drops Megatron's references to its NCCL groups, so every test's + # communicators (and their GPU buffers) piled up in the persistent workers (~20 groups and + # ~0.4 GiB per test), slowing later tests until they timed out. Destroy all but the default. + for group in list(torch.distributed.distributed_c10d._world.pg_map): + if group is not torch.distributed.GroupMember.WORLD: + torch.distributed.destroy_process_group(group) torch.cuda.empty_cache() @@ -90,16 +102,17 @@ def _make_pool(world_size): ) -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def _pool_cache(): - """Module-scoped cache of worker pools keyed by world_size. + """Session-scoped cache of worker pools keyed by world_size. Spinning up a pool cold-imports the full torch/megatron/modelopt stack per worker - (~tens of seconds), so fixtures that request the same world_size share one pool - instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and - ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at - module end, so workers are never reused across modules (avoids cross-test - state contamination). The ``dist_workers*`` fixtures are deliberately function-scoped so + (75-100 s in CI), so every module that requests the same world_size shares one pool, + e.g. on a 2-GPU runner ``dist_workers`` and ``dist_workers_size_2`` are both size 2. + Isolation between tests comes from ``megatron_worker_teardown`` (model-parallel state and + process groups after every job), from ``_reset_shared_pools`` (process-global torch settings + at each module boundary) and from ``get_live_pool`` replacing a pool that a hung, timed-out + or crashed job killed. The ``dist_workers*`` fixtures are deliberately function-scoped so ``get_live_pool`` can swap in a new pool after one was killed; they therefore cannot be requested from a module- or session-scoped fixture. """ @@ -109,21 +122,29 @@ def _pool_cache(): pool.shutdown() +@pytest.fixture(scope="module", autouse=True) +def _reset_shared_pools(_pool_cache): + """Workers outlive test modules: undo what the previous module did to process-global torch state.""" + for pool in _pool_cache.values(): + if not pool.broken: + pool.run(reset_worker_state) + + def _get_pool(cache, world_size): # A hung test kills its pool; a fresh one keeps a single hang from timing out every later - # test in the module (that cascade cost the nightly 20+ min). + # test (that cascade cost the nightly 20+ min). return get_live_pool(cache, world_size, lambda: _make_pool(world_size)) @pytest.fixture def dist_workers(_pool_cache): - """Pool (cached per module) with world_size=torch.cuda.device_count().""" + """Shared pool with world_size=torch.cuda.device_count().""" return _get_pool(_pool_cache, torch.cuda.device_count()) @pytest.fixture def dist_workers_size_1(_pool_cache): - """Pool (cached per module) with world_size=1 for tests that require a single process.""" + """Shared pool with world_size=1 for tests that require a single process.""" return _get_pool(_pool_cache, 1) From e0c6e212f54ff8350f9ee3ad1db26ff00e6e9ada Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:28:48 +0000 Subject: [PATCH 2/3] Keep a failed pool reset from erroring a module A pool that lives for the whole session turns one poisoned job into many errors. In the first CI run of this PR, test_lora_then_quantize hit the known Blackwell "illegal memory access" in an NCCL broadcast. A CUDA error of that kind is sticky, so the next jobs in the same workers failed too, and because _reset_shared_pools ran a job in the poisoned pool at the start of every module, every test of every later module errored, including the ones that use no pool (3 failed, 180 errors). With a pool per module such a flake cost one module. DistributedWorkerPool.run now retires a pool after any failed job (#2836), so get_live_pool hands the next test a new one and _reset_shared_pools skips a retired pool. What is left is a reset job that itself fails or finds a dead worker: _reset_shared_pools no longer lets that error the module. run() retires the pool, and the next test that needs one gets a new pool. Testing: the Megatron lane needs GPUs, so this PR's gpu_megatron job is the check. The retirement it relies on is covered on CPU by test_a_failed_job_retires_the_pool_even_when_every_rank_replied in tests/gpu/torch/utils/test_distributed_worker_pool.py. Signed-off-by: Shengliang Xu --- tests/gpu_megatron/conftest.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index c53a24abd49..0912fc39584 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib +import queue import pytest import torch @@ -111,8 +112,9 @@ def _pool_cache(): e.g. on a 2-GPU runner ``dist_workers`` and ``dist_workers_size_2`` are both size 2. Isolation between tests comes from ``megatron_worker_teardown`` (model-parallel state and process groups after every job), from ``_reset_shared_pools`` (process-global torch settings - at each module boundary) and from ``get_live_pool`` replacing a pool that a hung, timed-out - or crashed job killed. The ``dist_workers*`` fixtures are deliberately function-scoped so + at each module boundary) and from ``get_live_pool`` replacing a pool that a failed, hung, + timed-out or crashed job retired (a sticky CUDA error would otherwise fail every later + test). The ``dist_workers*`` fixtures are deliberately function-scoped so ``get_live_pool`` can swap in a new pool after one was killed; they therefore cannot be requested from a module- or session-scoped fixture. """ @@ -127,7 +129,10 @@ def _reset_shared_pools(_pool_cache): """Workers outlive test modules: undo what the previous module did to process-global torch state.""" for pool in _pool_cache.values(): if not pool.broken: - pool.run(reset_worker_state) + # A pool whose reset fails is retired by run(); the next test that needs one gets a new + # pool, and the module's tests, which may not use a pool at all, are not errored. + with contextlib.suppress(RuntimeError, queue.Empty): + pool.run(reset_worker_state) def _get_pool(cache, world_size): From a73bf1bd2cfad7af4e18437dda5ba75ce89489e4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 02:46:15 +0000 Subject: [PATCH 3/3] Collect garbage after destroying a test's process groups in the Megatron teardown The first two CI runs of this PR each hit "an illegal memory access was encountered" in test_megatron_peft (test_lora_then_quantize, then test_forward_with_one_lora[lora_config1]), after about 230 jobs in the shared worker. The lane almost never fails that way on other PRs, and it passed in the three runs of the draft this change was split from, whose teardown collected garbage right after destroying the groups; the version in this PR did not. The cause is not established. The groups are destroyed explicitly, but the finished test's Megatron and TE modules can still reference them, and without a collection their finalizers run at an arbitrary point of the next job, while its NCCL kernels may be in flight. Collecting right after the groups are destroyed, with nothing running, is the safe place for that. The cost is one collection per job, about 0.1-0.3 s. Testing: the Megatron lane needs GPUs, so this PR's gpu_megatron job is the check. The failure was intermittent, so the job has to pass repeatedly: it is re-run on this commit. Signed-off-by: Shengliang Xu --- tests/gpu_megatron/conftest.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 0912fc39584..1af63e694e5 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib +import gc import queue import pytest @@ -92,6 +93,9 @@ def megatron_worker_teardown(rank, world_size): for group in list(torch.distributed.distributed_c10d._world.pg_map): if group is not torch.distributed.GroupMember.WORLD: torch.distributed.destroy_process_group(group) + # The finished test's modules can still reference the destroyed groups. Collect them here, with + # no kernels in flight, rather than at an arbitrary point of the next job. + gc.collect() torch.cuda.empty_cache()