[OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules

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:

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()

tests/gpu_megatron/conftest.py

@@ -13,11 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
+import gc
+import queue
 
 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 +74,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 +87,15 @@ 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)
+    # 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()
 
 
@@ -90,16 +107,18 @@ 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 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.
     """
@@ -109,21 +128,32 @@ 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:
+            # 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):
     # 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)