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

[shengliangxu]

What does this PR do?

Type of change: CI (test runtime); changes tests/gpu_megatron and tests/_test_utils only

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). The loop is the one open PR CI follow-up: cache source-built wheels and extensions, bump CI containers, trim slow GPU tests #2836 adds for the same leak, comment included, so whichever PR lands second sees an identical hunk. 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; the workers record 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, loses a worker or fails on any rank retires the pool, and the next test gets a fresh one, for the cost of one pool start. DistributedWorkerPool.run retires a pool after a plain failed job because a CUDA error such as the known Blackwell illegal memory access is sticky: the workers stay poisoned, and with a session-wide pool every later test would fail with it. _reset_shared_pools also no longer lets a failed reset error the module's tests (they may not use a pool at all).

Measured saving on top of #2846's sysmon core: 216 s of the gpu_megatron lane (1,038 s to 822 s; see the time section). Without sysmon a pool start costs 76 s, and the saving would be about 17 min (an estimate; the series was not run that way). One pool start instead of 15 also means fewer chances to hit the ready timeout.

Time improvement (measured in CI)

gpu_megatron lane, pytest time:

main (#2845)	+ #2846	all slices up to #2854	+ this PR
2,240 s	1,212 s	1,038 s (1,029 s in #2852)	822 s

This PR on its own: 1,038 s -> 822 s, -216 s (-21%), against the two earlier runs (#2852: 1,029 s, #2854: 1,038 s). This PR's run was on a slightly slower runner (the heavy call phases of the tests it does not change took 6% longer than in #2854), so the gain is not runner luck. The setups of 12 s or more, which are the pool starts and the module-level compiles, fall from 16 (489 s) to 3 (246 s): the pool starts once instead of in each of 14 modules, and what remains is the extension build (124 s), the GDN module's Triton JIT (105 s) and one 18 s pool start.

With sysmon (#2846) a pool start takes 18-27 s instead of 75-77 s, so each avoided start saves less than the 76 s this PR was designed around. Without #2846 the saving would be 14 starts of 76 s plus the 3.6-10 s shutdowns, about 17 min; that is an estimate, because the series was never run without sysmon on this PR.

This is one CI run of the final head: the two runs before it failed on the sticky CUDA error and on the missing gc.collect() described below.

Method: every PR of the series has its own CI run, so "before" is the run of the PR(s) below that lack the change and "after" the runs that carry it. Times come from the job logs (--durations, or the gap between consecutive result lines). Where runner speed matters, the tests that no PR of the series touches scale each run, and the raw and the scaled medians are both shown. One test varies by 15-30% from run to run on its own, so smaller differences are not evidence.

Usage

Nothing user-facing (test fixtures).

Testing

• reset_worker_state has a gloo test next to the other pool tests (it runs on CPU): flags and seed changed in a job are back at their spawn-time values afterwards; it fails when the helper forgets one of the settings.
• The Megatron lane needs GPUs, so this PR's gpu_megatron job is the check: 822 s of pytest on the final head, against 1,038 s for the PR below (381 passed, 129 skipped, as before).
• The first CI run of this PR failed in gpu_megatron: the Blackwell illegal-memory-access flake hit test_lora_then_quantize and poisoned the shared pool, and every later module errored (3 failed, 180 errors), because the pool was reused after a failed job and _reset_shared_pools ran a job in it at each module start. The second commit retires a pool after any failed job (a gloo test on CPU covers it and fails without the change) and makes a failed reset harmless; with it the lane ended with 1 failure instead of 183, in test_forward_with_one_lora[lora_config1] (the same CUDA error). Re-running the job on that commit failed at the same test again, so it was not a random flake: the teardown of this PR, unlike the draft's, did not collect garbage after destroying the groups, and the draft passed the lane three times. The third commit adds that gc.collect(); the lane then passes. The cause is not established (the finished test's Megatron and TE modules can still reference the destroyed groups, and their finalizers then run at an arbitrary point of the next job). Open CI follow-up: cache source-built wheels and extensions, bump CI containers, trim slow GPU tests #2836 makes the same change to run() for its own reasons; the two are compatible, but its teardown loop has no collection, which a session-wide pool needs.
• Pre-commit passes.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: N/A (CI and tests)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: no new dependency
• Did you write any new necessary tests?: ✅ 1 test function in tests/gpu/torch/utils/test_distributed_worker_pool.py (gloo, CPU) for the new helper; the session pool itself is exercised by the whole gpu_megatron lane
• Did you update Changelog?: N/A (CI and tests)
• Did you get Claude approval on this PR?: N/A (optional, not run)

Additional Information

Slice 11 of 11 of the GPU-CI speed-up work, split out of the reference draft #2730 (not for merge; it holds the whole series with its measurements). Merge in order; each PR stands on its own and passes CI without the later ones.

1. [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist
2. [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI
3. [OMNIML-6179] [3/11] Run release-marked tests on nightly and on-demand GPU runs #2847 [OMNIML-6179] [3/11] Run release-marked tests on nightly and on-demand GPU runs
4. [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time #2848 [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time
5. [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models
6. [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data
7. [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests
8. [OMNIML-6180] [8/11] Trim the DFlash regression test #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test
9. [OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane
10. [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies #2854 [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies
11. [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules (this PR)

Summary by CodeRabbit

• Tests
  • Distributed GPU test workers now restore their initial PyTorch settings and random seed between test modules, with improved cleanup of GPU and Python resources.
  • Failed jobs now retire their worker pool even after all ranks respond, preventing broken workers from being reused.
  • Worker pools can be shared across test modules and are shut down at the end of the test session, improving test execution consistency.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 862e800c-7d4e-4e4d-acfb-bfaed566c828

📥 Commits

Reviewing files that changed from the base of the PR and between 1812f62 and 37ef132.

📒 Files selected for processing: 2

• tests/_test_utils/torch/distributed/utils.py
• tests/gpu/torch/utils/test_distributed_worker_pool.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 0 remain after this review.

📝 Walkthrough

Distributed worker utilities now capture and reset Torch state and retire pools after failed jobs. Megatron fixtures reuse pools across modules, reset state at module boundaries, and clean up CUDA and process-group resources during teardown.

Changes: Distributed worker lifecycle

Layer / File(s)	Summary
Capture and reset worker state tests/_test_utils/torch/distributed/utils.py, tests/gpu/torch/utils/test_distributed_worker_pool.py	Workers capture their initial Torch settings. reset_worker_state restores them, resets the seed to 1234, collects garbage, and clears the CUDA cache. Tests verify the reset behavior.
Retire pools after failed jobs tests/_test_utils/torch/distributed/utils.py, tests/gpu/torch/utils/test_distributed_worker_pool.py	DistributedWorkerPool.run aborts the pool before raising a worker error after a failed job. Tests verify pool retirement when all ranks reply.
Reuse and clean up Megatron pools tests/gpu_megatron/conftest.py	Megatron fixtures share pools across modules, reset healthy pools at module boundaries, and shut them down at session end. Worker teardown synchronizes CUDA, destroys non-default process groups, and collects garbage.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Other

Sequence Diagram(s)

sequenceDiagram
  participant Fixture as module-scoped reset fixture
  participant Pool as DistributedWorkerPool
  participant Worker as persistent worker
  Fixture->>Pool: run reset_worker_state
  Pool->>Worker: restore settings and reset seed
  Worker->>Worker: collect garbage and clear CUDA cache

Loading

Suggested reviewers: kevalmorabia97

Merge Risk: ⚪ Minimal · up to 37ef1

No concrete merge-blocking risk remains in the reviewed changes.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 40.74% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 27 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: sharing one Megatron worker pool across the gpu_megatron modules.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The pull request changes only three test files under tests/; it does not modify the modelopt package, examples, pyproject.toml, or requirements.txt. The added Python code contains no new…

Full details: Security Anti-Patterns

Explanation

PASS. The pull request changes only three test files under tests/; it does not modify the modelopt package, examples, pyproject.toml, or requirements.txt. The added Python code contains no new torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval, exec, or # nosec usage. No dependency changes are present.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[shengliangxu]

/ok to test 3c4611f

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.94%. Comparing base (229f069) to head (37ef132).

Additional details and impacted files

@@                         Coverage Diff                         @@
##           shengliangx/ci-10-pool-dead-worker    #2855   +/-   ##
===================================================================
  Coverage                               78.94%   78.94%           
===================================================================
  Files                                     658      658           
  Lines                                   74023    74023           
===================================================================
+ Hits                                    58435    58436    +1     
+ Misses                                  15588    15587    -1     

Flag	Coverage Δ
unit	60.50% <ø> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 9f2fca6

[shengliangxu]

/ok to test ca6126d

[shengliangxu]

/ok to test dee192a