[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). 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 has retired a pool after a plain failed job since CI follow-up: cache source-built wheels and extensions, bump CI containers, trim slow GPU tests #2836, which matters more with a session-wide pool: a CUDA error such as the known Blackwell illegal memory access is sticky, the workers stay poisoned, and every later test would fail with it. _reset_shared_pools 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	slices 1-10 (the earlier base)	+ 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 earlier base, [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies #2854's run (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 retired a pool after any failed job and made 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). CI follow-up: cache source-built wheels and extensions, bump CI containers, trim slow GPU tests #2836 has since landed the retirement in run() for its own reasons (with its own gloo test), so the second commit now only makes a failed reset harmless; main's teardown has neither the process-group loop nor the collection, and a session-wide pool needs both.
• 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). Slices 1, 2, 8 and 10 (#2845, #2846, #2852, #2854) have merged. Slice 5 (#2849) targets main and merges independently; the other open slices form a stack and merge bottom-up. Each PR stands on its own and passes CI without the later ones. The times above are from the series' CI runs before 2026-10-21, when all slices formed one linear stack; a PR number in a method note names that PR's run from then.

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 (merged)
2. [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI (merged)
3. [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs #2847 [OMNIML-6179] [3/11] Run manual-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 (targets main)
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 (merged)
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 (merged)
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.
  • Worker teardown now synchronizes CUDA, removes non-default process groups, and cleans up Python and GPU resources.
  • Worker pools are shared across test modules and shut down at the end of the test session for more consistent test execution.

[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]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 339ca31
• Latest commit, not reviewed: cf0d7de

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

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: af0ee2d6-9197-4475-839b-a4e664414448

📥 Commits

Reviewing files that changed from the base of the PR and between 39d221e and 339ca31.

📒 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; 7 remain after this review.

📝 Walkthrough

Distributed worker utilities now capture and reset Torch state. Megatron fixtures reuse pools across modules, reset healthy workers 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.
Reuse and clean up Megatron pools tests/gpu_megatron/conftest.py	Megatron fixtures share pools across modules, reset healthy pools at module boundaries, and clean up CUDA and process-group resources during worker teardown.

Priority: ⬇️ Low

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

Change: Other

Sequence Diagram(s)

sequenceDiagram
  participant Fixture as module 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 339ca

No actionable merge-blocking risk is established.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

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

Passed checks: 5 passed

Check name	Status	Explanation
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 reviewed range changes only three files under tests/; it does not change any Python file under modelopt/ or examples/, nor pyproject.toml or requirements.txt. No added `torch.load(…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The reviewed range changes only three files under tests/; it does not change any Python file under modelopt/ or examples/, nor pyproject.toml or requirements.txt. No added torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), trust_remote_code=True, eval, exec, or # nosec pattern was found.

✨ 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 70.08%. Comparing base (94f2f23) to head (cf0d7de).

Additional details and impacted files

@@                        Coverage Diff                        @@
##           shengliangx/ci-9-torch-onnx-lane    #2855   +/-   ##
=================================================================
  Coverage                             70.08%   70.08%           
=================================================================
  Files                                   659      659           
  Lines                                 74081    74081           
=================================================================
+ Hits                                  51917    51918    +1     
+ Misses                                22164    22163    -1     

Flag	Coverage Δ
unit	60.51% <ø> (+<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