[OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time by shengliangxu · Pull Request #2848 · NVIDIA/Model-Optimizer

shengliangxu · 2026-10-17T01:41:29Z

What does this PR do?

Type of change: CI (test reliability)

DistributedWorkerPool waited 120 s for its workers to import the Megatron/TE/modelopt stack and report ready. The start takes 75-77 s in 14 modules of the PR run I read and has a p90 of 94-98 s, so slow hosts time out and the whole module errors: 15 jobs (903 job-minutes) in two weeks, 15 of the 63 failed gpu_megatron PR jobs in that period. It waits 300 s.

An earlier version of this PR also gave test_mamba_search_space more time. Main's #2836 has since set its cap to 150 s, and that holds: in 58 gpu_megatron jobs from the last six days the test took 64-132 s (median 73 s, none above 150 s), so nothing is left to change there.

Time improvement (measured in CI)

None in a passing run: this PR raises one limit that only failing runs reach.

In the 9 gpu_megatron runs of this series, the worker-pool start took 75-77 s before [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 and 18-27 s after it, against the old 120 s ready timeout. No run came near the limit, so the series cannot show the failures this PR prevents (the sysmon core of [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 also takes most starts out of reach of the old ready timeout).
What it prevents, from the CI history of the two weeks before the series: the ready timeout failed 15 jobs (903 job-minutes), 15 of the 63 failed gpu_megatron PR jobs. Each failed job costs the author a re-run of the lane, 20-45 min of job time in this series.

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.

Testing

Pre-commit passes. Not run on a GPU here: the effect is on cold-cache runs, so this PR's gpu_megatron job can only show that nothing regresses.

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?: N/A (a timeout in the shared test utilities)
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 4 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.

[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
[OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI
[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
[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 (this PR)
[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
[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
[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
[OMNIML-6180] [8/11] Trim the DFlash regression test #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test
[OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane
[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
[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

Summary by CodeRabbit

Bug Fixes
- Increased the wait time for distributed worker initialization to improve reliability during startup.

copy-pr-bot · 2026-10-17T01:41:33Z

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 · 2026-10-17T01:41:38Z

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: 22a4971

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: 02e38d3f-eac2-4f39-b840-0683f96553a7

📥 Commits

Reviewing files that changed from the base of the PR and between 3fa0758 and 22a4971.

📒 Files selected for processing: 1

tests/_test_utils/torch/distributed/utils.py

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

📝 Walkthrough

DistributedWorkerPool.__init__ now waits up to 300 seconds for each worker’s "ready" message, up from 120 seconds.

Changes: Worker readiness timeout

Layer / File(s)	Summary
Worker readiness timeout `tests/_test_utils/torch/distributed/utils.py`	The readiness wait increased from 120 to 300 seconds. Worker startup and readiness validation remain unchanged.

Priority: ➖ Normal

Estimated code review effort: 1 (Trivial) | ~5 minutes

Change: Bug fix

Suggested reviewers: kevalmorabia97

Merge Risk: ⚪ Minimal · up to 22a49

This CI-only change extends worker startup tolerance without introducing an actionable merge risk.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 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 describes the main change: increasing the timeout for slow gpu_megatron cold starts. It is specific and consistent with the pull request changes.
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 `tests/_test_utils/torch/distributed/utils.py`, outside the modelopt package and examples. The change raises a worker-ready timeout from 120 to 300 seconds and adds…

Full details: Security Anti-Patterns

Explanation

PASS. The pull request changes only tests/_test_utils/torch/distributed/utils.py, outside the modelopt package and examples. The change raises a worker-ready timeout from 120 to 300 seconds and adds a timing comment. It adds no security-sensitive loading, evaluation, # nosec, or dependency patterns covered by this check.

✨ 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.}

codecov · 2026-10-17T01:47:56Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (shengliangx/ci-3-release-tests@9e1afc3). Learn more about missing BASE report.

Additional details and impacted files

@@                        Coverage Diff                        @@
##             shengliangx/ci-3-release-tests    #2848   +/-   ##
=================================================================
  Coverage                                  ?   76.90%           
=================================================================
  Files                                     ?      658           
  Lines                                     ?    74029           
  Branches                                  ?        0           
=================================================================
  Hits                                      ?    56932           
  Misses                                    ?    17097           
  Partials                                  ?        0

Flag	Coverage Δ
unit	`60.46% <ø> (?)`

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 · 2026-10-17T02:37:56Z

/ok to test 14186a5

shengliangxu · 2026-10-17T07:31:26Z

/ok to test be2bdd8

kevalmorabia97 · 2026-10-19T07:59:47Z

    assert not any(named_dynamic_modules(model))


+@pytest.mark.timeout(240)  # The first Mamba (Triton) kernel compile takes 75-100+ s on a cold cache


Partly handled in #2836
Increased from 120 -> 150s timeout as that seemed sufficient

Right, and this PR no longer touches it: it is rebased onto #2836, and the Mamba hunk is gone from the current head (22a2fcc). Your 150 s holds: in 58 gpu_megatron jobs from the last six days the test took 64-132 s (median 73 s), none above 150 s. What remains of this PR is the worker-pool ready timeout, 120 s to 300 s.

DistributedWorkerPool waited 120 s for its workers to import the Megatron/TE/ modelopt stack and report ready. The start takes 75-77 s in 14 modules of one PR run and has a p90 of 94-98 s, so slow hosts time out and the whole module errors: 15 jobs (903 job-minutes) in two weeks. Wait 300 s. Testing: ruff check and format with the pinned ruff 0.15.20; py_compile. Not run on a GPU runner. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…#2845) ### What does this PR do? Type of change: CI (unit-test runtime) The unit `linux` job gates every GPU and example workflow (their pr-gate jobs wait for it), and its tests are CPU-bound but run serially. This runs them under pytest-xdist on the runner's four vCPUs. The `linux` job and the multi-version matrix get `PYTEST_ADDOPTS=-n 4 --dist loadfile` and `OMP_NUM_THREADS=1`: one worker per vCPU (`-n auto` counts physical cores), whole files per worker, and one thread per worker, because with torch's default threads the runner is oversubscribed and slower than serial (1,481 s and 41 failures in my local runs). What else changes so the suite passes under xdist: - `noxfile.py`: the second pytest process that runs the `manual` tests clears `PYTEST_ADDOPTS`, so it stays serial. Those tests leak global state into each other, and a worker would run them in some order. - `tests/_test_utils/torch/quantization/onnx_export.py`: `TEST_MODELS` was a set of classes, whose iteration order differs between processes, so workers collected different test ids. It is a tuple now. - `test_hf_dflash2.py::test_overfits_a_single_batch` takes about 30 s on one thread on main (53 s before #2836 halved its loop to 30 steps; that version hit the 60 s default call timeout in the nightly on main on 2026-10-20). Its tiny model is built in bf16, and a CPU without native bf16 support emulates bf16 matmuls (about 0.9 s per step, 0.1 s in float32). The model is float32 now, which takes the test to about 4 s on one thread, so it needs no timeout override; the loop and both assertions are unchanged. - `test_full_algorithms_vision.py::test_search_constraints` assigned `ConstraintsFunc._get_true_latency` and never restored it. Serial runs never noticed, because `test_evaluate_constraints.py` sorts first. Under xdist a worker can get the files in the other order, and the Python 3.10 and 3.14 legs failed two tests after it. It uses `monkeypatch` now. - `tests/conftest.py`: the unit per-test cap is 120 s on every platform. `test_dist.py::test_data_parallel` takes 8-15 s serial and 25-35 s with four workers sharing four vCPUs, and went over 60 s on the slowest runner. A hung test is still reported, after 2 min instead of 1. Why: on current main the `linux` job has a median of 12.9 min (27 runs of other PRs since #2823 merged), and a pull request's GPU run waits a median 14.5 min for it (125 runs of 33 PRs). ### Time improvement (measured in CI) Unit `linux` job: serial on current main (27 runs of other PRs since #2823 merged) against the runs of this series with `-n 4`. | | before (serial) | with this PR | |---|---|---| | pytest | 630 s median (355-709 s) | 479 s in this PR's own run; 294 s median over the 11 runs of the series (166-479 s) | | whole job | 12.9 min median (8.8-14.8) | 10.9 min in this PR's own run; 7.0 min median over the 11 runs (5.0-10.9) | | unit workflow: `linux`, then the multi-version and Windows legs | 24.1 min median (12 runs) | 17.2 min median (the 7 runs that ran every leg) | - The runs of #2846 and above also carry the sysmon core of #2846 (about 22% off the unit suite on its own, measured locally), so the 294 s median is xdist plus sysmon; this PR's own run is the xdist-only number. It was the slowest of the 11 (166-479 s for nearly the same tests: runner variance), so it understates the gain. Locally on 4 CPUs with coverage as in CI, the base took 825 s serial and 332-344 s with this configuration (3 runs). - The GPU, Example and Regression workflows hold their jobs until the unit `linux` check is done (`pr-gate / wait-checks`, at least 5 min). That wait is a median 14.5 min (IQR 12.9-15.8) on 125 GPU runs of 33 PRs since #2823; in this PR's own GPU run it was 9.7 min (its `linux` job took 10.9 min). With a 7 min `linux` job instead of 12.9 min, a GPU run requested right after a push starts about 6 min earlier. That figure is derived from the job times: the later runs of the series were requested after their unit check had finished (the wait sat at its 5.2 min floor in 7 of the 9 GPU runs), so they cannot show it. 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. Locally: `OMP_NUM_THREADS=1 pytest tests/unit -n 4 --dist loadfile`. ### Testing - Full `tests/unit` on 4 CPUs with coverage as in CI, on the base before #2823: serial 825 s, this configuration 332-344 s in three repeated runs (4,724 passed, 0 failed). - This PR's own CI run: the `linux` job took 10.9 min (479 s of pytest) against the 12.9 min serial median; see the time section above for the later runs of the series. - `tests/unit` without `puzzletron` (it needs `hydra`, absent from my local environment) in reversed and in shuffled file order shows no other order dependence; `test_search_constraints` followed by `test_evaluate_constraints.py` in one process failed before the fix and passes after it. - `test_hf_dflash2.py` (33 tests) on one CPU thread: 13.5 s, of which `test_overfits_a_single_batch` takes about 4 s (30 s with the bf16 model on main; the next slowest test takes about 1.3 s). - Pre-commit (all hooks) passes. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (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 (`pytest-xdist` is already in the `dev-test` extra on main) - Did you write any new necessary tests?: N/A (two existing tests fixed or sped up so they pass under xdist) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (CI and tests) - Did you get Claude approval on this PR?: N/A (optional, not run) ### Additional Information Slice 1 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. #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist **(this PR)** 2. #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI 3. #2847 [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs 4. #2848 [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time 5. #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models 6. #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data 7. #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests 8. #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test 9. #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane 10. #2854 [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies 11. #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules  ## Summary by CodeRabbit * **Tests** * Improved consistency and performance of automated test runs across platforms and Python configurations through controlled parallel execution. * Preserved predictable model-check ordering and isolated temporary test changes. * Increased unit-test time limits to better support process-spawning and slower runs. * Improved test reliability by explicitly controlling numeric precision and ensuring serial execution for manual CPU tests. * No end-user functionality or product behavior changed.  --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu added this pull request to stack #2856 October 17, 2026 01:41

This was referenced Oct 17, 2026

[OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist #2845

Merged

[OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846

Open

shengliangxu changed the title ~~shengliangx/ci 4 megatron cold start~~ [4/11] Give gpu_megatron's slow cold starts more time Oct 17, 2026

shengliangxu mentioned this pull request Oct 17, 2026

[OMNIML-6182] [1/2] Run the required checks in GitHub's merge queue #2858

Draft

shengliangxu changed the title ~~[4/11] Give gpu_megatron's slow cold starts more time~~ [4/12] Give gpu_megatron's slow cold starts more time Oct 17, 2026

shengliangxu removed this pull request from stack #2856 October 17, 2026 06:58

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch from 14186a5 to be2bdd8 Compare October 17, 2026 07:02

shengliangxu added this pull request to stack #2859 October 17, 2026 07:02

shengliangxu changed the title ~~[4/12] Give gpu_megatron's slow cold starts more time~~ [4/11] Give gpu_megatron's slow cold starts more time Oct 17, 2026

shengliangxu changed the title ~~[4/11] Give gpu_megatron's slow cold starts more time~~ [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time Oct 18, 2026

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch from be2bdd8 to 202efca Compare October 18, 2026 23:15

shengliangxu marked this pull request as ready for review October 19, 2026 03:44

shengliangxu requested a review from a team as a code owner October 19, 2026 03:44

coderabbitai Bot approved these changes Oct 19, 2026

View reviewed changes

kevalmorabia97 reviewed Oct 19, 2026

View reviewed changes

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch from 202efca to 14d1c95 Compare October 19, 2026 21:34

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch from 14d1c95 to 891f950 Compare October 20, 2026 01:16

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch 2 times, most recently from 3fa0758 to 22a2fcc Compare October 20, 2026 18:58

shengliangxu force-pushed the shengliangx/ci-4-megatron-cold-start branch from 22a2fcc to 22a4971 Compare October 20, 2026 19:43

Navigation Menu

Repository navigation

[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#2848
shengliangxu wants to merge 1 commit into
shengliangx/ci-3-release-testsfrom
shengliangx/ci-4-megatron-cold-start

shengliangxu commented Oct 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

copy-pr-bot Bot commented Oct 17, 2026

Uh oh!

coderabbitai Bot commented Oct 17, 2026 •

edited

Loading

Reviews paused

Uh oh!

codecov Bot commented Oct 17, 2026 •

edited

Loading

Uh oh!

shengliangxu commented Oct 17, 2026

Uh oh!

shengliangxu commented Oct 17, 2026

Uh oh!

kevalmorabia97 Oct 19, 2026

Uh oh!

shengliangxu Oct 20, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		assert not any(named_dynamic_modules(model))


		@pytest.mark.timeout(240) # The first Mamba (Triton) kernel compile takes 75-100+ s on a cold cache

Conversation

shengliangxu commented Oct 17, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Time improvement (measured in CI)

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Oct 17, 2026

Uh oh!

coderabbitai Bot commented Oct 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Uh oh!

codecov Bot commented Oct 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

shengliangxu commented Oct 17, 2026

Uh oh!

shengliangxu commented Oct 17, 2026

Uh oh!

kevalmorabia97 Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shengliangxu commented Oct 17, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Oct 17, 2026 •

edited

Loading

codecov Bot commented Oct 17, 2026 •

edited

Loading