diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index ffe731c0dc4..b4e656434e8 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: @@ -179,6 +210,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 36dce35968f..ea55e942800 100644 --- a/tests/gpu/torch/utils/test_distributed_worker_pool.py +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -19,8 +19,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): @@ -41,6 +42,32 @@ def _teardown_barrier(rank, size): dist.barrier() # like megatron_worker_teardown: blocks while a peer is still in the job +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: @@ -83,3 +110,11 @@ def test_a_failed_job_retires_the_pool_even_when_every_rank_replied(): assert pool.broken finally: pool.shutdown() + + +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..1af63e694e5 100644 --- a/tests/gpu_megatron/conftest.py +++ b/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)