Repository navigation
[OMNIML-6177] [Draft, not for merge] GPU CI speed-up series (runtime measurement) - #2730
Closed
shengliangxu wants to merge 34 commits into
Closed
shengliangxu wants to merge 34 commits into
shengliangxu wants to merge 34 commits into
Conversation
|
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. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Comment |
Collaborator
Author
|
/ok to test 5064ef3 |
Contributor
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2730 +/- ##
==========================================
+ Coverage 76.89% 80.67% +3.77%
==========================================
Files 658 646 -12
Lines 73649 70522 -3127
==========================================
+ Hits 56636 56895 +259
+ Misses 17013 13627 -3386
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
shengliangxu
added a commit
that referenced
this pull request
Oct 10, 2026
The Windows unit job failed at collection: the parametrize decorator of
test_worker_death_fails_the_step_not_the_session evaluates -signal.SIGKILL when the module is
imported, and signal has no SIGKILL on Windows. The module's skip_on_windows fixture only runs
at test setup, so the AttributeError aborted the whole session ("Interrupted: 1 error during
collection") and no Windows test ran at all (PR #2730, Unit tests / windows).
Take the parameter from a module-level getattr(signal, "SIGKILL", 9). The script text that sends
the signal still names signal.SIGKILL; it only runs where the test runs, which is POSIX.
Testing: collecting the committed version with signal.SIGKILL, SIGHUP, SIGUSR1 and os.setsid,
killpg and getpgid removed reproduces the CI error. The fixed module, test_distributed_worker_pool.py
and test_example_run_command.py collect (41 tests) under the same conditions. Not run on Windows.
The collection error hid every other Windows result, so the job needs a re-run to show what else,
if anything, fails there.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 10, 2026
test_search_constraints replaced ConstraintsFunc._get_true_latency with a fake by plain assignment and never put the original back, so the fake stayed in place for the rest of the process. test_evaluate_constraints.py::test_evaluate_latency_constraints and ::test_percent_limits use the real method through a fake runtime client (latency 100 and 200) and fail when they run after it: the leaked fake returns 280.5 for the tiny resnet, so the 20% latency limit comes out as 56.1 instead of 40, and a 200 limit is not met. Serial runs never saw it, because test_evaluate_constraints.py sorts before test_full_algorithms_vision.py. With pytest-xdist (-n 4 --dist loadfile) a worker can get the files in the other order. In PR #2730 the same worker ran test_search_constraints and then the two victims in the Python 3.10 and 3.14 legs, which failed; in the legs where they landed on different workers, or in the other order, they passed. Use monkeypatch, so the patch is undone after the test. Testing: running test_search_constraints[False-True] and then test_evaluate_constraints.py in one process failed both tests with the CI values before the change (56.10 against 40.0, and assert False is True) and passes after it, and ConstraintsFunc._get_true_latency is the original function again afterwards. The whole unit suite in reversed and in shuffled file order is run next to look for other leaks of this kind. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 10, 2026
tests/unit/torch/quantization/test_dist.py::test_data_parallel hit the 60 s default cap in the Python 3.12 / torch 2.13 leg of PR #2730. The test spawns two processes and takes 8-15 s in serial runs. With pytest-xdist it takes 25-35 s in every other leg, and 60 s or more on the slowest runner: four workers and the child processes of the multi-process tests (test_dist, the three test_*data_parallel* autoquant tests, the worker-pool tests) share the four vCPUs of a CI runner. A hung test is still reported, after 2 min instead of 1. Windows already had 120 s. Testing: the durations above come from the job logs of the PR run and of the serial run before it. Not re-run on a runner: the cap only adds headroom. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
/ok to test 524d784 |
The unit workflow's `linux` job gates every GPU and example workflow: their pr-gate jobs wait for it (poseidon/wait-for-status-checks, `^linux$`). That wait is a median 13.8 min, 23% of a PR's GPU run, and the job itself takes a median 14.8 min. The tests are CPU-bound and run serially, so use the runner's four vCPUs. - unit_tests.yml: `PYTEST_ADDOPTS=-n 4 --dist loadfile` on the `linux` job and the multi-version matrix. One worker per vCPU (`-n auto` counts physical cores), whole files per worker, and `OMP_NUM_THREADS=1`: with torch's default threads per worker the runner is oversubscribed and the run is slower than serial (1,481 s and 41 failures in my local runs). - 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 xdist workers collected different test ids. Use a tuple. - test_hf_dflash2.py: test_overfits_a_single_batch needs more than the default 60 s call timeout when one thread runs it (about 53 s with torch's default threads). Give it 300 s. pytest-xdist is already in the dev-test extra on main (#2771); uv.lock is refreshed by the weekly bump_uv_lock workflow. Testing: the full tests/unit suite, on 4 CPUs with coverage as in CI, on the base before the unit-test cleanup of #2823. Serial 825 s; this configuration 332-344 s in three repeated runs (4,724 passed, 0 failed). The other matrix entries (Python 3.10, 3.13, 3.14, torch 2.8, transformers min) take 4.5-5.3 min. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
pytest-cov builds the term-missing, html and xml reports at the end of the session (inside its runtestloop wrapper), over about 640 modelopt files. CI uploads only coverage.xml to Codecov, so the other two cost time and a 600-line table in every log for nothing. The cost shows up as ~45 s on the last test of every lane (gpu 46 s, gpu_megatron 49 s, TRT-LLM 43.5 s, vLLM 43.8 s, hf_ptq, llm_qat, llm_eval) and as a 71 s gap before the last result line of the unit `linux` job, which every GPU and example workflow waits for. It also pushes jobs over their timeouts: 59 of 143 timed-out PR jobs in two weeks had already finished 95% or more of their tests and died in the coverage tail or the post steps. Timed on this tree: html 20.2 s, xml 6.2 s, term-missing 4.6 s. A 17-test run takes 50.9 s with the old addopts and 25.1 s with this one. Developers who want the other reports pass --cov-report=html or --cov-report=term-missing. Testing: ran pytest with --cov before and after on a small test selection and checked that only coverage.xml is written and the run is shorter. Not yet run on a real CI runner. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Python 3.12 defaults to coverage's C tracer, which slows every line of Python that runs under it. Every lane starts coverage in each spawned worker and subprocess through COVERAGE_PROCESS_START, so the import-heavy start-ups pay for it: the 76 s Megatron worker-pool start, the 22 s hf_ptq processes of the unified-export tests, the puzzletron spawns, the vLLM engine core. COVERAGE_CORE=sysmon records the same lines through sys.monitoring at a fraction of the overhead. Measured (the first three here, the last by #2727 on its hf_ptq example lane): - serial tests/unit run on 4 CPUs: 572 s without coverage, 825 s with the C tracer, 640 s with sysmon; - `import modelopt.torch.quantization.extensions` with this repo's coverage config: 17.1 s native, 23.6 s with the C tracer, 18.3 s with sysmon, and the child process of a spawn writes its data file too, with the same 362 files and 13,821 executed lines; - one TinyLlama fp8 plus TRT-LLM deploy test: 64.8 s without coverage, 103.3 s with the C tracer, 70.5 s with sysmon, identical covered lines. Applies to the unit `linux` job (Python 3.12), gpu_tests.yml and regression_tests.yml. The example runner is left to #2727, which makes the same change there. sysmon needs Python 3.12 or newer; on an older interpreter coverage warns and keeps its default core, so a lane that turns out to run an older Python is slower, not broken. The gpu and gpu_megatron logs show 3.12. Testing: actionlint is clean on all workflows. Not yet run on a real CI runner; the first run should be compared against the previous coverage.xml. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…kernels The extension cache hardly ever helps, and nothing caches Triton at all. Extension cache. Its key was `rtxpro6000-<container image>`, shared by every job on an image, and `actions/cache` saves only when the job succeeds: - gpu_megatron and the megatron_bridge example job share nemo:26.08. The example job, which builds fewer extensions, saved a 0.9 MB entry first. The gpu_megatron log then says "Cache hit ... 920,006 B ... not saving cache", so it compiles the CUDA extensions in every job and never saves them: test_cuda_ext has a 125 s setup in 92% of 334 PR jobs (median 132.7 s). - main is never seeded for the PR branches: the 2-GPU nightly gpu job is cancelled at its 60-minute limit every night, and a cancelled job saves nothing. In the gpu lane 40-51% of PR jobs miss the cache (test_cuda_ext median 0.1 s, mean 18 s, p90 43 s); hf_ptq shows the same 25.7 s extension JIT in 18% of jobs. - Triton-only edits under modelopt/torch/kernels/quantization invalidated the CUDA entry, because every file there was hashed. Triton. Every container starts with an empty Triton cache, so each job recompiles the same kernels: the GDN fixture 103 s (about 30 compiles), the Mamba first calls 137 s, the FLA fixture 67 s, the NVFP4 FP8-sweep kernel about 93 s, Qwen3.5 vision +24.6 s, the first vLLM 0.30 engine 102 s against 35 s warm. Sibling parametrizations with the same shapes take 1-3 s. The open #2572 measured the Mamba compile at 97 s cold and 16.6 s with a warm cache. Changes: - cache-extensions now restores two entries (it no longer saves, because the composite action cannot run a step after the job): CUDA extensions under an exact key hashed over the .cu/.cuh/.cpp/.h sources, the conv loader and the two Python loaders only, and the Triton cache (TRITON_CACHE_DIR) under a key over modelopt/torch/kernels/**/*.py with a restore-keys prefix. Triton entries are content-addressed, so restoring an older one is harmless. - New cache-extensions/save action, run with `if: always()` after the tests in gpu_tests.yml and _example_tests_runner.yml. It saves what was not an exact hit, also when the job failed or timed out. - The cache key now includes the lane (`rtxpro6000-<lane>-<image>`). Seeding main still needs the nightly to get through (separate change); until then each PR branch saves its own entries on its first run. Testing: actionlint is clean on all workflows, and with a deliberately broken input or output name it reports the error, so the wiring of the new inputs and outputs is checked. Assumes actions/cache/restore@v6 and /save@v6 exist next to actions/cache@v6. Not yet run on a real CI runner: confirm the restored cache sizes and that a second job on the same lane hits. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Two timeouts in gpu_megatron fail PR jobs for reasons that are not bugs: - test_mamba_search_space has a 74 s median call (the first Mamba/Triton kernel compile on a cold cache) against the 120 s default cap of tests/gpu_megatron. Of 330 PR jobs that reached its module, 38 (11.5%) have no result line for it and 5 more error out; it failed in 24 runs over two weeks (1,425 job-minutes), and runs missing its result line end as failures 16 times in 38 against 15 in 292. The next test then absorbs 207-235 s (the rest of the setup plus the cap). Give it 240 s, like the other compile-heavy test in test_megatron.py. - 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. Wait 300 s. Together they are 34 of the 63 failed gpu_megatron PR jobs in that period. Testing: ruff check and format with the pinned ruff 0.15.20; py_compile. Not run on a GPU runner. The open #2572 splits the Mamba compile into its own test, which is compatible with this. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
main's DistributedWorkerPool (#2832) already kills its workers and marks itself broken when a job is interrupted or a rank fails and its peers hang, and the fixtures replace a broken pool through get_live_pool. Two gaps remain, and the shared Megatron pool of a following change needs two helpers. - A worker that exits (an illegal memory access aborts the process, an OOM kill) was only noticed when the 600 s result timeout ran out, and a start that failed waited out the ready timeout and left its workers behind. The pool now checks the workers' liveness once a second while it waits: a dead worker fails the run or the start within a second, the pool is marked broken and the other workers are killed with their stacks printed, as for any other abort. - ready_timeout is a constructor argument (300 s by default, as set by the earlier timeout change). - destroy_extra_process_groups() destroys every group but the default one. Libraries that create groups per model drop their references but not the groups, and each NCCL group keeps a communicator and a watchdog thread. - reset_worker_state() restores the torch settings a fresh worker starts with (cudnn and TF32 flags, deterministic algorithms, RNG), which set_seed() and similar helpers change. Workers record them right after they start. Testing: tests/unit/test_distributed_worker_pool.py runs a two-process gloo pool on CPU: a job error keeps the workers; an interrupt raised by a signal handler, as pytest-timeout does, kills them and breaks the pool; a worker that exits fails the run in under 20 s and kills its peer; a start that fails raises quickly and leaves no process behind; the group cleanup and the flag reset work. 6 tests in 30 s. main's tests/gpu/torch/utils/test_distributed_worker_pool.py (3 tests, gloo) passes against the merged pool, run on CPU. ruff 0.15.20 is clean. Not covered here, and still needing GPU CI: NCCL, CUDA poisoning, real Megatron jobs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
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. main's fixtures (#2832) are function-scoped and get their
pool from get_live_pool, which replaces a pool that a hang killed; both stay as
they are. 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) and runs gc;
- _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, in the pools that already exist. It is not part of _get_pool, which
runs once per test and would add a job and a gc to each;
- get_live_pool and the pool's abort handling (main's, plus the previous change
for a dead worker): a job that times out, is interrupted or loses a worker kills
the pool and the next test gets a fresh one, for the cost of one pool start. That
also covers CUDA poisoning after the known Blackwell illegal-memory-access flake.
Expected saving: about 17 min of the gpu_megatron job (36 min of tests to about
19), and one pool start instead of 15 means fewer chances to hit the ready
timeout.
Testing: ruff 0.15.20. The pool, group-cleanup and flag-reset behaviour is
covered by the CPU tests of the previous change. The Megatron lane needs GPUs: in
draft PR #2730, before this rebase onto #2832, the gpu_megatron pytest step took
862 s on cold caches and 336 s on warm ones, against 2,022 s for a baseline run
(this change plus the JIT and Triton caches). The conflict resolution with main's
function-scoped fixtures and get_live_pool has not run on a GPU.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every gpu-lane job compiles three CUDA packages against the container's torch before its first test: fast-hadamard-transform (171-172 s), mamba_ssm (185 s on the nightly, 278 s on a PR runner) and causal-conv1d (5 s). That is 6-7.6 min of the 31-min PR job; the session starts pytest 8.7 min after its first command (PR run 37703177445, nightly run 37708492798). No other pip wheel build in the 67 job logs I read took more than 20 s. A new `gpu_wheels` nox session builds the three wheels with `pip wheel --no-build-isolation --no-deps` into $WHEELHOUSE, and the `gpu` session installs the wheel files found there instead of compiling. Without WHEELHOUSE, or with an empty one, `gpu` runs exactly the pip commands it ran before. In the workflow, for the gpu leg only: - the wheelhouse is restored under wheelhouse-<image>-<ISO week>-<hash of noxfile.py>. The week keeps the "latest release" semantics of the unpinned mamba_ssm, causal-conv1d and fast-hadamard-transform HEAD within a week; the hash drops the entry when a requirement changes, so a stale wheel is never installed for a changed pin; - only runs on main (the nightly, or a manual run) build and save. The build is its own step, so an entry is saved only when all three wheels exist and a failed build saves nothing. Runs on pull-request branches only restore: a PR entry (~600 MB: 50 + 343 + 212 MB) would count against the repo's cache budget and no other PR could read it; - a miss (the first runs of a new week until the nightly has saved, or a PR whose noxfile.py differs from main's) compiles in the session as today. The step runs after "Setup environment variables" so the build sees the same environment the in-session build did. Expected saving: most of the 6-8 min on every gpu-lane job that hits, less the restore and install of ~600 MB (not measured). Testing: ruff 0.15.20 and actionlint (a wrong step id is caught). The noxfile was run against a fake session: without WHEELHOUSE and with an empty one it issues the old pip commands, with wheels present it installs them by path and falls back per package, and gpu_wheels refuses to run without WHEELHOUSE. `pip wheel --no-build-isolation --no-deps -w <new dir>` on dummy projects created the directory and wheels named causal_conv1d-*, fast_hadamard_transform-* and mamba_ssm-*, which install and import. The real CUDA builds, the cache sizes and the restore time have not been run: the first nightly after merging does the build and save, the next run shows the hit. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Tests marked `release` collect only when pytest gets --run-release, and no workflow passes it, so a test marked `release` never runs in CI (the docs say the same of the example release suite). Set PYTEST_ADDOPTS=--run-release in the "Run gpu tests" step for every run that is not a pull-request run: the nightly schedule and manual dispatches. That lets slow or network-heavy checks leave the pull-request jobs without leaving CI. Pull-request runs are unchanged. Nothing under tests/gpu* is marked release yet; the next commit marks the first test. The release example suites (hf_ptq/test_deploy.py, gpt-oss QAT) belong to example_tests.yml, which does not get the flag, so the docs sentence that said no workflow passes it now says no example workflow does. Testing: actionlint is clean. A stub test marked `release`, collected under the repo's tests/conftest.py, is skipped with "Only run when --run-release is given" when PYTEST_ADDOPTS is empty and runs when it is --run-release. The expression is `!startsWith(github.ref, 'refs/heads/pull-request/') && '--run-release' || ''`, which gives '' for a pull-request ref and '--run-release' otherwise. Not run on GitHub. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_megatron_generate_and_mmlu[tp] builds a 28-layer Qwen3-0.6B Megatron model, downloads 1.2-1.5 GB of real weights into a fresh container cache and scores about 1,400 MMLU prompts: a 32 s call in every pull-request gpu_megatron job. It is the only test in the lane that downloads real weights. Mark it `release` so pytest skips it unless --run-release is given. The importer shape and name tests and the generate control-flow tests in the same module keep running on pull requests. This test is also the only one that checks real-weight accuracy (0.36 < MMLU < 0.39), so a numeric regression in the import or generation path is no longer caught on a pull request. The previous commit passes --run-release on the nightly and on-demand runs, so the real-weight accuracy check still runs there, a day after a change instead of on the pull request. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran a stub module with the same decorators under the repo's tests/conftest.py on CPU: without the flag the four parametrized ids are skipped with "Only run when --run-release is given", with it they all run. Not run: the test itself, because this machine has no GPU and its Megatron is too old to import the module, and the flag on GitHub (the previous commit checked it with a stub test only). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The gpu_vllm lane boots a tiny vLLM engine per architecture. A 64-token request with max_num_seqs=1 needs only a few KV blocks, but the defaults size the KV pool from gpu_memory_utilization (millions of block objects) and the scheduler token budget to thousands of tokens. Pass num_gpu_blocks_override=32 and max_num_batched_tokens=64 to the Qwen3-MoE, DeepSeek-V3 and DeepSeek-V3 fp8 engines. The sparse, GLM, DeepSeek-V4 and GPT-OSS engines (other KV layouts) and the dense Llama engine are unchanged. Estimated saving: 1-3 s per engine boot. The engine fixtures are function-scoped on main, so that is four boots per vLLM version: 4-12 s. Not done: the Llama engine could go (about 28 s per vLLM version: the Qwen3-MoE and DeepSeek engines cover its linear types, and DeepSeek's shared expert has the merged and row-parallel ones), but test_tiny_llama_quantize now also hosts the amax and recipe reload checks of #2627, which would have to move first. Testing: ruff check and ruff format with the pinned ruff 0.15.20. Read the vLLM 0.20.0 and 0.30.0 sources: EngineArgs accepts both arguments, kv_cache_utils honours the override, and the scheduler rules (max_num_batched_tokens >= max_num_seqs, and >= max_model_len without chunked prefill) hold at max_model_len=64 and max_num_seqs=1. The same settings on these three engines passed the gpu_vllm jobs of both vLLM versions in draft PR #2730, on the base before #2627; the rebased file has not run on a GPU. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_dflash_training trained 3 epochs (1,500 steps) and test_dflash_resume then trained a fourth epoch (500 more steps). On the 1-GPU pull-request runner the recorded call durations solve to about 33 s of start-up plus 55-73 ms per step, so the two tests spend roughly 115-140 s and 60-70 s. The loss assertion only needs a final loss below 3.0, and the L40 baseline is already 1.78 at step 500. Train for 500 steps and resume for 20 more (max_steps 520). The resume still goes through the whole restore-and-train path; only its tail is shorter. It had no assertion besides the exit status, so it now reads trainer_state.json and checks that global_step is 520 and that a log entry follows step 500 (the training summary the resumed run appends). The step count is fixed with max_steps instead of num_train_epochs because the regression job runs on 2 GPUs off pull requests, and `accelerate launch` starts one process per visible GPU when there is no accelerate config. An epoch is then 250 steps, so num_train_epochs=1 would train 250 steps with the last loss logged at step 200, which may not clear the unchanged `< 3.0` bound, and 1.04 epochs would resume for 10 steps, not 20. logging_steps and the loss assertions (final < first, final < 3.0) are unchanged. Only the comment on the final-loss bound now cites the step-500 baseline. Estimated saving: about 80 s per pull-request regression job (55 s in training, 25 s in the resume) and about 30 s on the 2-GPU job. Risk: the 1.78 baseline at step 500 comes from a 1,500-step run. With max_steps=500 the linear learning-rate schedule now decays to zero at step 500, so the final loss can differ from that baseline. The bound leaves 1.2 of margin on it. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Read the transformers 4.57.1 sources (max_steps overrides num_train_epochs; with the steps save strategy a checkpoint is saved at step 500 and again at the end of training) and the accelerate 1.10.0 launcher defaults. The recipe's training schema accepts extra keys, and launcher examples already pass training.max_steps. Not run: pytest, because this machine has no GPU and the test downloads Qwen3-0.6B. GPU CI has to confirm that the final loss stays below 3.0 after 500 steps and that the resume run ends at global_step 520. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The megatron_bridge conftest resets Megatron's global state before and after every test: destroy_model_parallel, destroy_rerun_state_machine, then gc.collect() and torch.cuda.empty_cache(). With torch, Megatron and modelopt imported the collection walks a large heap and costs about 1.3 s per test in the recorded durations. 77 tests in this directory are pure Python (the mlflow wiring tests and the checkpoint-path tests), never run an example step and leave nothing to reset, so the lane pays about 100 s per job for nothing. The megatron_bridge job runs at 68-75 minutes against its 75-minute nightly limit. Reset the Megatron state only for tests that request `num_gpus`, which is how every test that runs an example step reaches a GPU in this directory. The environment restore stays for every test: it is cheap, and the NVTE_* leak it guards against does not depend on a GPU. Tests that run steps keep both resets, and run_example_in_process and run_torchrun_in_process already reset at the start of each step. A future test that runs a step without taking `num_gpus` would lose the reset around it, though its steps would still reset at their start. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran a throwaway pytest module (not part of this change) on CPU against the patched conftest with the real fixture, counting calls to reset_megatron_global_state: plain tests trigger none, tests that request `num_gpus` directly or through another fixture trigger one before and one after, and the environment is restored in both cases. The same module fails against the old conftest. Not run: the megatron_bridge suite itself, because megatron.bridge and the GPUs are not available on this machine. The next GPU CI run should show the setup and teardown time of the mlflow tests drop to near zero. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Six tests in test_megatron_preprocess_data.py tokenize with gpt2 from the Hugging Face Hub. megatron_preprocess_data loads the tokenizer once for the vocabulary size and again in every worker, so each test makes several hub round trips, and none of them depends on gpt2's vocabulary: they check output naming, non-empty .bin and .idx files, and that max_tokens shortens the output. Switch them to the tiny_tokenizer_path fixture, the checked-in byte-level tokenizer that the blend test already uses. Three of the six read only local data (the jsonl path, the jsonl max_tokens and the gzip tests), so the tokenizer download was their only network access. The PR history lists the jsonl max_tokens test, with the HF split max_tokens test, among the lane's 120 s timeouts (3 failures each); it no longer has a network dependency to time out on. The other three (HF split max_tokens, HF resume and the HF streaming warning) still read the real HF dataset, which is the code path they exist to cover; only their tokenizer changes. The Qwen3 tests keep their real tokenizer: they need its chat template and EOS handling. The blend test already uses the tiny tokenizer, and its HF split source is the path it covers, so it is unchanged. No test ids, assertions or markers change. Measured on a CPU host with a cold Hugging Face cache, the nine tests in the module take 38.5 s before and 24.4 s after: jsonl path 3.3 s to 0.01 s, jsonl max_tokens 3.6 s to 0.13 s, gzip 1.6 s to 0.07 s, HF split max_tokens 7.0 s to 4.9 s, HF resume 5.0 s to 3.1 s, HF streaming warning 4.4 s to 2.8 s. CI jobs start without a warm hub cache, so a similar saving is expected. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran the module on CPU, next to a copy of the tiny_tokenizer_path fixture from tests/gpu_megatron/conftest.py, with the Hub reachable: 4 passed and 5 xpassed before and after, and the collected test ids are identical. Not run: the gpu_megatron lane itself, which needs GPUs and prebuilds CUDA extensions in its conftest. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The three ONNX test files added with #2678 (test_dla_qdq_rules.py, test_dla_scalar_constants.py and test_quantization_initializer_names.py, 68 test cases) quantize with calibration_eps=["cpu"] and run the result on the CPUExecutionProvider, so they need no GPU. They lived under tests/gpu/onnx and ran in the `gpu` job, where they cost 26.8 s per job (1.5 + 5.3 + 20.0 s in the CI logs). The `gpu` job is the second-longest lane of a PR run; the unit job runs the same tests in a few seconds on its xdist workers. Move them to tests/unit/onnx/quantization/ without changing their content. Testing: ran the three files on CPU with CUDA_VISIBLE_DEVICES empty and the repository's own modelopt on the path. 68 passed in 7.1 s at the old location and 68 passed in 13.3 s from the new one (this includes the unit conftest's session setup) on Python 3.12 with onnx 1.21.0 and onnxruntime 1.24.4. On Python 3.10.19 with onnxruntime 1.22.0, the oldest runtime pyproject allows, 68 passed in 15.3 s. The whole tests/unit/onnx directory still passes (756 passed, 12 xfailed). Not run: the Windows unit job and the GPU lane; CI should show the `gpu` job losing about 27 s. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_concat_elim_int8 and test_concat_elim_fp8 take about 22 s each in the CI logs, and their assertions only look at where the concat_elimination pass places Q/DQ nodes. quantize() calibrated a (1, 128, 240, 64) three-conv model on the default [cpu, cuda:0, trt] providers. It now calibrates a 32-channel 16x16 model on the CPU. The Convs stay above the 16-channel limit that graph_selection.py applies to INT8 and FP8, and both modes produce the same 8 Q and 8 DQ nodes as with the default size. build_conv_concat_model takes optional channels, height and width arguments. Their defaults are the old values, and with a seeded RNG the default model serializes to the same bytes as before. Testing: ruff check and format with the pinned ruff 0.15.20. With the repository's own modelopt on CPU, the two tests pass in 1.3 s; quantizing the default and the small model in both modes gives Q=8, DQ=8 each time. Not run: the gpu job. From the CI timings this should save about 40 s per job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_benchmark.py built the 64x32x224x224 autotune model for every integration test. For the Python-API benchmark that means filling 103 M random floats and copying about 1.2 GB per iteration, and the five tests that benchmark a model take 31 s in the PR run (7.9, 7.9, 5.9, 4.8 and 4.8 s in run 37703177445). None of them looks at the latency beyond "a positive float". The fixtures now build a 1x32x32x32 model. The two trtexec tests also pass --duration=0: trtexec measures for at least 3 s by default whatever the iteration count, and these tests ask for two iterations. _create_simple_conv_onnx_model takes optional batch_size and spatial_size arguments with the old values as defaults, so the unit autotuner tests that share it build the same model (their 161 tests pass). Testing: ruff check and format with the pinned ruff 0.15.20, and the unit autotune tests with the repository's own modelopt on CPU. Not run: the TensorRT tests, because there is no TensorRT here. Estimated saving is about 20 s per gpu job; the CI run will show it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_export_quantized_model[True] takes 44 s in the PR run (run 37703177445), mostly inside trtexec. trtexec measures for at least 3 s per run by default whatever --iterations says, and the workflow builds and measures the baseline plus a number of Q/DQ schemes with it. The test only needs a latency to rank the schemes by, so pass --duration=0 and let trtexec stop after the 100 requested iterations. That is still 100 iterations per measurement; only the extra iterations the soak added are gone. The Python-API parameter ignores trtexec_args. The model keeps its size: the test asserts that the best scheme carries Q/DQ nodes, which needs a model whose latency is well above launch noise, and timing_runs stays at 100 for the same reason. Testing: ruff check and format with the pinned ruff 0.15.20. Not run: the test, because there is no TensorRT here. Expected saving is 3 s per measured scheme, probably 10-20 s per gpu job; the CI run will show it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…other - save_dummy_dataset wrote 2500 copies of one conversation into each of the train and valid splits. The tiny test tokenizer turns a conversation into 1424 tokens, so that is 1738 blocks of 2048 tokens, while the pruning scores read 100 blocks and the scoring and realize steps 2 each. Every dataloader build tokenizes the whole split, and a run builds three of them. The 2500 came from an older 8192-token Llama-3 setup. Use 400 conversations (278 blocks). On a CPU the three builds plus the dataset save took 2.8 s per test with 2500 conversations and 0.36 s with 400. The scoring step checks lm_loss with a tolerance of 0.15; on random-init toy Llamas the validation loss moved by at most 0.02 between the two sizes, because the same text repeats and only the choice of blocks changes. - calc_subblock_stats ran 12 passes (batch sizes 64, 96 and 128 times four dtype sets), but the MIP and the teacher memory and parameter checks only read the entry for the batch_size in mip.subblock_stats_args, which is 96 in all nine configs. batch_sizes: [96] gives 4 passes (0.56 s to 0.18 s for a toy Llama on a CPU; MoE configs have more subblocks per pass). - test_nas_convert_ffn_pruning ran the same mtn.convert call on the same toy Llama and config as test_nas_search (21.6 s in the CI logs, with its own 2-rank spawn). Its two assertions move into test_nas_search, which checks them with the other rank-0 assertions after the search. Left alone: the Nemotron-3 toy keeps its 128 experts, because shrinking it changes EXPECTED_TEACHER_MEMORY_MIB and EXPECTED_TEACHER_NUM_PARAMS, which need a GPU runner with mamba_ssm to regenerate. Testing: timed the dataloader builds, the subblock-stats passes and the loss sensitivity on CPU with the repository's own helpers and dataloader code, collected tests/gpu/torch/puzzletron (16 tests, one fewer than before), and ran ruff. Not run: the puzzletron tests themselves, which need GPUs and mamba_ssm. Expect the gpu job to lose roughly 50 s (the removed test is 22 s of it); the CI run will show the real number. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…or batch probing Each of the 11 test_unified_hf_export_and_check_safetensors params starts a fresh hf_ptq process, 22 s per param in the CI logs. The test only checks the structure of the exported safetensors, so three parts of that run are not needed: - Calibration text came from cnn_dailymail on the Hub, a network round trip per param. Use a local jsonl of 64 random-word texts instead, written by a new helper, write_calib_jsonl, in _test_utils/torch/calib_data.py. Each text fills the 512-token calibration window and differs from the others, so the activation statistics stay varied. - hf_ptq probed for the largest calibration batch. --batch_size 64 puts the whole calibration set (calib_size 64) in one batch. - hf_ptq ran model.generate() before and after quantization. --skip_generate skips those previews; the export is not affected. Testing: no GPU here, so I ran the test on a CPU host through a python shim that lets hf_ptq pass its GPU guard and passes --device cpu. 8 of the 11 params pass with the new command line, scale-equality assertions included: fp8, nvfp4, int4_awq, w4a8_awq_beta, int8_weight_only, w4a16_nvfp4, the Qwen3 MoE nvfp4 and the GPT-OSS fp8. The other three (nvfp4 AWQ-lite, SVDQuant and the MSE sweep) need CUDA or Triton and did not run. The dataloader returns one (64, 512) batch with a full attention mask from the helper's file. Not run: the GPU job itself. From the per-param CI timings this should save about 38 s per `gpu` job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_gptq_e2e_flow streamed cnn_dailymail from the Hub for each of its three params, a network round trip per param (about 6 s per job in the CI logs) and a test that can fail on a Hub hiccup. The GPTQ calibration only needs some text, and the test has no accuracy assertion, so it now reads 8 texts from the local jsonl that write_calib_jsonl (added with the unified-export change) creates. Testing: ruff check and format with the pinned ruff 0.15.20. On a CPU host the GPTQ flow completes on that jsonl for the FP8 and INT4 blockwise configs; the NVFP4 config needs CUDA. Not run: the GPU test itself. Expected saving is a few seconds per gpu job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Two cuts to the torch_onnx example lane. swinv2_tiny runs all five formats on top of swin_tiny, which already covers every format through the same exporter and TensorRT path; swinv2 only adds its cosine attention and post-norm blocks. On PRs (GITHUB_REF under refs/heads/pull-request/) skip swinv2 int8, mxfp8 and auto and keep fp8 and nvfp4, which still exercise those blocks. Scheduled and local runs keep every combination, and test ids are unchanged. swinv2_tiny[auto] alone averages 121 s over 304 CI runs (the 22nd slowest example test); the three together cost about 110 s of every PR run. build_trt_engine let trtexec run its default benchmark after the build (200 ms warm-up, then at least 3 s and 10 iterations), and the output is discarded on success. Run a single iteration instead (--warmUp=0 --duration=0 --iterations=1): the engine still executes once, but about 3 s of timing per engine goes away, roughly 40 s per lane run. --trt_build users see no difference besides it returning sooner. Testing: ruff 0.15.20 (check and format), py_compile and pytest --collect-only on CPU (24 tests, ids unchanged). With the example command stubbed, GITHUB_REF=refs/heads/pull-request/123 skips exactly swinv2 int8, mxfp8 and auto, and any other ref skips none. The two in-process mapping tests that import the script pass on CPU, and build_trt_engine passes the three flags to a fake trtexec. Not run: the lane itself (needs a GPU, TensorRT and model downloads) and the flags against a real trtexec, so GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every example step is a fresh `python script.py` that spends 13-22 s of CI time importing torch, transformers and modelopt before it does any work, and an example test runs several steps: 41 launches in the diffusers lane, 21 in torch_onnx, 11 in llm_qat. Only megatron_bridge avoids them today, with a runner that executes in the pytest process and needs the get_args()/main() convention. ExampleRunner runs `python <script>.py ...` steps in one long-lived spawned worker. A step runs as python would run it (runpy as __main__, sys.argv, the example dir as cwd and first on sys.path, output captured at the fd level), and afterwards everything it changed is put back: argv, cwd, sys.path, os.environ, signal handlers, logging handlers and levels, warning filters, RNGs, torch backend flags, the multiprocessing start method, and the modules imported from examples/, so the `utils` or `config` of different examples do not collide. Third-party and modelopt imports stay cached, which is the saving: on a CPU box a step that imports timm, transformers and modelopt.torch.quantization takes 9.4 s as a subprocess, and 0.35 s in the worker after the first one (10.2 s). Per-lane hooks wrap steps for registries the generic restore cannot know about. The worker is a separate process, not the pytest process. A hung step is killed with its children and fails with the output it had printed, and the next step gets a fresh worker. The same happens when pytest-timeout interrupts the wait. A step that takes its worker down fails alone, the worker is recycled after any failed step, and pytest holds no CUDA context. close() lets the worker exit so that coverage writes its data. run_example_command hands an installed runner the steps it accepts (an allowlist of scripts) together with the step's environment; every other step runs as before. StepFailed subclasses CalledProcessError, so callers see what a subprocess would raise, and the Hub retry loop reads its captured output. Plain callables such as the megatron_bridge runner keep their contract. MODELOPT_EXAMPLE_RUNNER=subprocess runs every step as a subprocess, to bisect a suspected leak between steps; =inprocess runs them in the calling process so that pdb works. No lane installs the runner yet. Testing: tests/unit/test_example_runner.py (32 CPU tests, about 10 s) and tests/unit/test_example_run_command.py pass, also under pytest-xdist (-n 4 --dist loadfile) and under COVERAGE_PROCESS_START with COVERAGE_CORE=sysmon (about 30 s). Dropping each of 19 behaviours of the runner (every restore, the process-group kill, the daemon flag, the pipe close, the graceful close) fails at least one test. ruff 0.15.20 and typos are clean. Not run: any GPU lane, so memory released after a step and CUDA RNG state are unvalidated and GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The lane starts `python torch_quant_to_onnx.py` 19 times per PR run, and every start spends 15-20 s importing torch, timm, datasets and modelopt before the step quantizes a one-sample model. Install an ExampleRunner for the lane in tests/examples/torch_onnx/conftest.py, per test as the megatron_bridge conftest does, so that one session-wide worker pays the imports once. The plan estimates -140 to -400 s per PR run. test_hf_embedding_quant_to_onnx.py is not on the runner's list and keeps running as a subprocess, and every step does when MODELOPT_EXAMPLE_RUNNER=subprocess. What the script changes in the process, and what puts it back: sys.path gets ../onnx_ptq, which is restored with the rest of sys.path; the sibling modules evaluation and download_example_onnx are imported from examples/ and dropped after each step; the multiprocessing start method is no longer forced (previous commits) and would be restored anyway; the seeds and cudnn flags of evaluation.py are restored with the other RNG and backend state. The ResNet recipe also adds `_add_resnet_residual_quantizers` to modelopt's CUSTOM_POST_CONVERSION_PLUGINS, a registry that would then apply to every later step. The new keep_post_conversion_plugins hook restores it after each step. Testing: ruff 0.15.20 and typos are clean, and tests/examples/torch_onnx still collects 24 tests on CPU. With the conftest active and CUDA hidden, test_torch_onnx_recipe_flag and test_torch_onnx_auto_quantize_recipe pass through the worker (11.7 s for the first step, 1.3 s for the second) and so do the two mapping tests. The real script through the worker gives the same ONNX op histogram as a subprocess for vit_tiny int8 and fp8: 12.9 s for the first step and 0.6 s for each later one, against 11.7 s per subprocess. A ResNet fp8 recipe step followed by a probe shows the plugin registry restored with the hook and leaked without it, and a unit test of the hook fails when its restore is removed. Not run: anything on a GPU, TensorRT or the --trt_build steps (trtexec runs as a child of the worker), so GPU memory across steps in one process, the worker's coverage data and the lane's timings are unvalidated and GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The Windows unit job failed at collection: the parametrize decorator of
test_worker_death_fails_the_step_not_the_session evaluates -signal.SIGKILL when the module is
imported, and signal has no SIGKILL on Windows. The module's skip_on_windows fixture only runs
at test setup, so the AttributeError aborted the whole session ("Interrupted: 1 error during
collection") and no Windows test ran at all (PR #2730, Unit tests / windows).
Take the parameter from a module-level getattr(signal, "SIGKILL", 9). The script text that sends
the signal still names signal.SIGKILL; it only runs where the test runs, which is POSIX.
Testing: collecting the committed version with signal.SIGKILL, SIGHUP, SIGUSR1 and os.setsid,
killpg and getpgid removed reproduces the CI error. The fixed module, test_distributed_worker_pool.py
and test_example_run_command.py collect (41 tests) under the same conditions. Not run on Windows.
The collection error hid every other Windows result, so the job needs a re-run to show what else,
if anything, fails there.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_search_constraints replaced ConstraintsFunc._get_true_latency with a fake by plain assignment and never put the original back, so the fake stayed in place for the rest of the process. test_evaluate_constraints.py::test_evaluate_latency_constraints and ::test_percent_limits use the real method through a fake runtime client (latency 100 and 200) and fail when they run after it: the leaked fake returns 280.5 for the tiny resnet, so the 20% latency limit comes out as 56.1 instead of 40, and a 200 limit is not met. Serial runs never saw it, because test_evaluate_constraints.py sorts before test_full_algorithms_vision.py. With pytest-xdist (-n 4 --dist loadfile) a worker can get the files in the other order. In PR #2730 the same worker ran test_search_constraints and then the two victims in the Python 3.10 and 3.14 legs, which failed; in the legs where they landed on different workers, or in the other order, they passed. Use monkeypatch, so the patch is undone after the test. Testing: running test_search_constraints[False-True] and then test_evaluate_constraints.py in one process failed both tests with the CI values before the change (56.10 against 40.0, and assert False is True) and passes after it, and ConstraintsFunc._get_true_latency is the original function again afterwards. The whole unit suite in reversed and in shuffled file order is run next to look for other leaks of this kind. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
tests/unit/torch/quantization/test_dist.py::test_data_parallel hit the 60 s default cap in the Python 3.12 / torch 2.13 leg of PR #2730. The test spawns two processes and takes 8-15 s in serial runs. With pytest-xdist it takes 25-35 s in every other leg, and 60 s or more on the slowest runner: four workers and the child processes of the multi-process tests (test_dist, the three test_*data_parallel* autoquant tests, the worker-pool tests) share the four vCPUs of a CI runner. A hung test is still reported, after 2 min instead of 1. Windows already had 120 s. Testing: the durations above come from the job logs of the PR run and of the serial run before it. Not re-run on a runner: the cap only adds headroom. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/ci-speed
branch
from
October 17, 2026 00:38
524d784 to
954ac75
Compare
Collaborator
Author
|
/ok to test 954ac75 |
This was referenced Oct 17, 2026
[OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models
#2849
Merged
shengliangxu
changed the title
shengliangxu
added a commit
that referenced
this pull request
Oct 20, 2026
…#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 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 21, 2026
### What does this PR do? Type of change: CI (test runtime) Sets `COVERAGE_CORE=sysmon` on the unit `linux` job, `gpu_tests.yml` and `regression_tests.yml`. Python 3.12 defaults to coverage's C tracer, which slows every line of Python run under it, and every lane starts coverage in each spawned worker and subprocess through `COVERAGE_PROCESS_START`, so import-heavy start-ups pay for it (the 76 s Megatron pool start, the 22 s hf_ptq processes of the unified-export tests, the puzzletron spawns, the vLLM engine core). sysmon records the same lines through `sys.monitoring`. The example runner is left to #2727, which makes the same change there. On a Python older than 3.12 coverage warns and keeps its default core, so a lane that turned out to run an older Python would be slower, not broken (the gpu and gpu_megatron logs show 3.12). ### Time improvement (measured in CI) pytest time of each lane in the run of #2845 (the base of this PR) and in this PR's run, one run each: | lane | #2845 | this PR | change | |---|---|---|---| | gpu | 2,093 s | 1,289 s | -38% | | gpu_megatron | 2,240 s | 1,212 s | -46% | | gpu_vllm v0.20.0 | 204 s | 170 s | -17% | | gpu_vllm v0.30.0 | 426 s | 328 s | -23% | | regression | 366 s | 268 s | -27% | | unit `linux` (xdist) | 479 s | 305 s | -36% | Workflow wall time: GPU tests 59.2 -> 39.6 min, Regression 21.3 -> 15.2 min. - **How much of that is this PR.** Two single runs of the same code differ by 15-30% on their own: the eight example lanes that run over 100 s, which this PR does not touch, moved by -4% to -31% (llm_sparsity +38%) between the same two runs. The deterministic part is the Megatron worker-pool start, the setup of the first test of each gpu_megatron module: 75-77 s in each of 14 modules before (all setups of 12 s or more add up to 1,365 s) and 22-27 s after (575 s), which is 790 s of the lane's 1,028 s. Locally the unit suite takes 825 s with the C tracer, 640 s with sysmon and 572 s without coverage (4 CPUs, serial). - **A `coverage.xml`-only report change was dropped from this PR.** The runs above are of its earlier head, which also had it. It showed no CI effect: the pytest summary line came 1-4 s after the last test in the #2845 runs, which still wrote the html report, and 0-3 s in this PR's runs (median 1.1 s over the 82 lane runs of the series), so the table is the sysmon core's. It would only have saved developers' local runs (about 25 s per pytest session). 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 Measured on this tree (tests/unit on 4 CPUs serial: 572 s without coverage, 825 s with the C tracer, 640 s with sysmon; `import modelopt.torch.quantization.extensions`: 17.1 s native, 23.6 s with the C tracer, 18.3 s with sysmon, with the same 362 files and 13,821 executed lines in the spawned child's data file; one TinyLlama fp8 plus TRT-LLM deploy test: 64.8 s without coverage, 103.3 s with the C tracer, 70.5 s with sysmon, identical covered lines). Pre-commit and actionlint pass. The CI runs of the series are the check on real runners (time section above). ### 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 - Did you write any new necessary tests?: N/A (CI configuration) - 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 2 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 5, 8 and 10 (#2849, #2852, #2854) target main and merge 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. 1. #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist (merged) 2. #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI **(this PR)** 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 (targets main) 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 (targets main) 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 (targets main) 11. #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Updated coverage collection for GPU, regression, and Linux unit-test runs to use Python’s monitoring-based coverage core on supported Python 3.12+ environments. * This configuration reduces overhead during coverage-enabled test execution, particularly for import-heavy worker processes. * No changes to application functionality or public interfaces are included. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 21, 2026
### What does this PR do? Type of change: CI (test runtime, tests only) **The DFlash regression trains 500 steps and resumes for 20.** `test_dflash_training` trained 3 epochs (1,500 steps) and `test_dflash_resume` then a fourth epoch (500 more). On the 1-GPU PR runner the recorded durations solve to about 33 s of start-up plus 55-73 ms per step, so the two tests spend roughly 115-140 s and 60-70 s. The loss assertion only needs a final loss below 3.0, and the L40 baseline is already 1.78 at step 500. The resume still goes through the whole restore-and-train path, and it now reads `trainer_state.json` and checks that `global_step` is 520 and that a log entry follows step 500 (it had no assertion besides the exit status). The step count is fixed with `max_steps` rather than `num_train_epochs`, because the regression job runs on 2 GPUs off pull requests, where `accelerate launch` starts one process per visible GPU and an epoch is 250 steps. Measured: 81 s per PR regression job. Risk: the 1.78 baseline comes from a 1,500-step run, and with `max_steps=500` the linear schedule decays to zero at step 500, so the final loss can differ from it; the bound leaves 1.2 of margin. ### Time improvement (measured in CI) | lane | | before | after | change | |---|---|---|---|---| | regression | `test_dflash_training` | 118 s | 62 s | -56 s | | regression | `test_dflash_resume` | 56 s | 31 s | -25 s | | regression | **lane pytest** | 268 s | 170 s | **-98 s (-37%)** | DFlash: before = #2846, #2848 and #2850 (#2849 ran on a slower runner), after = #2852, #2854 and #2855 (median; the lane took 168 and 170 s on the normal-speed runners and 243 s on the slow one). On its own, on main (2026-10-21): this PR's regression job against #2854's, which has the same base and lacks this change. The unchanged DFlash tests took the same time within about 3% (equally fast runners). `test_dflash_training` takes 144 s before and 83 s after, `test_dflash_resume` 73 s and 47 s, and the lane's pytest time goes from 374 s to 291 s (-83 s). The lane passes (5 passed), so the loss stays below 3.0 after 500 steps and the resume ends at `global_step` 520. A vLLM change that sized the KV pool of the Qwen3-MoE and DeepSeek engines down was dropped from this PR. The runs above are of its earlier head, which had it: the engines it touched booted 5-6 s per vLLM version faster than the ones it left alone, in line with the 1-3 s per boot expected and below the boot-to-boot noise (the unchanged Llama engine took 84-99 s on v0.30.0), so it saved nothing measurable. 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 I read the transformers 4.57.1 sources (`max_steps` overrides `num_train_epochs`; with the steps save strategy a checkpoint is saved at step 500 and again at the end) and the accelerate 1.10.0 launcher defaults. Not run on a GPU here: this PR's regression job is the check, in particular that the final DFlash loss stays below 3.0 after 500 steps and that the resume ends at `global_step` 520. Pre-commit 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 - Did you write any new necessary tests?: N/A (existing tests made cheaper; the DFlash resume test now also asserts on its result) - 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 8 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 5, 8 and 10 (#2849, #2852, #2854) target main and merge 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. #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist (merged) 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 (targets main) 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 (targets main) **(this PR)** 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 (targets main) 11. #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Updated regression coverage to use fixed training limits of 500 steps, with resumed runs extending to step 520. * Improved checkpoint and resume validation to confirm the final training step and expected post-resume progress logs. * Updated loss expectations to reflect the result at the 500-step training milestone. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 21, 2026
… a worker dies (#2854) ### What does this PR do? Type of change: CI (test reliability); changes `tests/_test_utils` only A worker that exits without reporting (an illegal memory access aborts the process, or an OOM kill) was only noticed when the result timeout ran out: 600 s for a job, 120 s for the start. Until then the parent waits on a queue nobody writes to, and the other workers wait in the job, typically in a collective, for their dead peer. A start that fails, for example on a bad backend, also waited out the whole ready timeout and left the workers it had spawned behind. `DistributedWorkerPool` now checks the workers' liveness once a second while it waits. A dead worker is noticed within a second and fails the run, or the start, naming its rank and exit code (about 4 s after the death, which includes the 3 s the pool already waits for the other workers' stack dumps). If another rank had already reported an exception, its traceback is kept in that error. The pool is marked `broken` and the other workers are killed with their stacks printed, the way a hung job or a failed rank already ends in main (`get_live_pool` then hands out a new pool). A start that fails kills the workers it had spawned. `ready_timeout` is a constructor argument (120 s by default, as before), so a test can reach the start timeout without waiting for it. This builds on the hang handling of #2832 (`broken`, `get_live_pool`, the stack dump) and does not change it. ### Time improvement (measured locally, not in CI) No lane time changes in a passing run, and no worker died in this series' CI runs, so CI cannot show it. What changes is how long a dead worker costs. CPU, gloo, two ranks, the pool test's own job (rank 0 exits with code 13, rank 1 keeps waiting), seconds from `pool.run(...)` (or the constructor) to the exception, on main and on this PR: | scenario | before | after | |---|---|---| | a rank dies mid-job, its peer is stuck in a wait, `job_timeout` 30 s (as in the test) | 33.1 s, `queue.Empty` | 4.0 s, `Worker process(es) died (rank, exit code): [(0, 13)]` | | the same with the default `job_timeout` of 600 s | 603.1 s, `queue.Empty` | 4.0 s | | the same, the peer waits in `dist.barrier()` | 33.3 s (the 30 s `failure_grace`) | 4.3 s | | the start fails (`backend="not-a-backend"`), default `ready_timeout` of 120 s | 120.1 s, `queue.Empty` | 5.0 s | | rank 0 dies at the start, rank 1 hangs in the rendezvous | 120.0 s, one live worker left behind | 5.0 s, no worker left | The after time is the 1 s liveness poll plus the 3 s the pool waits for the other workers' stack dumps before it kills them. At the production timeouts (120 s for the start, 600 s for a job) a dead worker therefore costs 4-5 s instead of 2-10 minutes of a CI job: the test fails at once and the next one gets a fresh pool from `get_live_pool`, instead of waiting for a queue nobody writes to. Cost: the four new cases add about 27 s to the `gpu` lane: the pool test module takes 43 s in this PR's run and 17 s in #2849's run, which has the same base (main) and lacks them (both from the job logs). Method: a small script (not part of this PR) builds the pool the way the repo's test does and times each scenario with `time.monotonic()` around the call, once per commit, on a CPU host. ### Usage ```python pool = DistributedWorkerPool(world_size=2, backend="nccl", ready_timeout=300) # new argument, default 120 ``` ### Testing `tests/gpu/torch/utils/test_distributed_worker_pool.py` runs the pool with gloo on CPU (it extends main's three tests; the new cases take 24 s together): - a rank that exits with code 13 while its peer waits fails the run in about 5 s and kills the peer. The test sets `job_timeout = 30`, so a regression fails there with `queue.Empty` instead of at the test's own cap. With the liveness check disabled it fails exactly that way. - the same test with a second job: rank 0 raises an exception and rank 1 exits with code 13 two seconds later, after rank 0 has reported. The error carries the exit code and rank 0's traceback. Without the change it carries only the exit code, and the case fails. - a start that fails raises and leaves no worker, both when every worker dies (a bad backend) and when the wait times out while they are still starting (`ready_timeout=0.01`). With `_abort()` removed from the failed-start path, the second case fails ("workers left behind"); the first passes either way, since all its workers die on their own, so it covers the liveness wiring only. Main's three tests (as reworked by #2836) still pass, and the three new cases fail against main's unmodified pool. Pre-commit passes. This PR is rebased onto #2836, which also edited the pool (reporting a failed rank before its teardown, retiring a pool after any failed job); the changes are in different places. The default start timeout stays 120 s here; #2848 raises it to 300 s, so whichever of the two merges second resolves that one line. ### 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 - Did you write any new necessary tests?: ✅ 2 test functions (4 cases) in `tests/gpu/torch/utils/test_distributed_worker_pool.py`, run on CPU with gloo; each fails without its fix - 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 10 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 5, 8 and 10 (#2849, #2852, #2854) target main and merge 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. #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist (merged) 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 (targets main) 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 (targets main) 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 (targets main) **(this PR)** 11. #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Distributed worker pools now detect workers that exit during startup or job execution and report failures instead of waiting for the full job timeout. * Worker startup supports a configurable readiness timeout with prompt timeout handling. * Failed or timed-out initialization now cleans up spawned worker processes, preventing orphaned processes and leaving the pool in a consistent state. * Job result retrieval now responds promptly when a worker fails. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 21, 2026
…e unit lane, small models (#2849) ### What does this PR do? Type of change: CI (test runtime, tests only) Four changes to the ONNX tests of the `gpu` job: - **Three CPU-only files move to the unit lane** (`tests/gpu/onnx/quantization/` to `tests/unit/onnx/quantization/`, content unchanged): `test_dla_qdq_rules.py`, `test_dla_scalar_constants.py` and `test_quantization_initializer_names.py` (added with #2678, 68 cases). They quantize with `calibration_eps=["cpu"]` and run the result on the CPUExecutionProvider, so they need no GPU, but they cost 26.8 s per `gpu` job (1.5 + 5.3 + 20.0 s in the CI logs); the unit job runs them in a few seconds on its xdist workers. - **`test_concat_elim` calibrates a 32-channel 16x16 model on the CPU** instead of the (1, 128, 240, 64) three-conv model on `[cpu, cuda:0, trt]`: the two tests took about 22 s each on the base (15 s together in the series' CI runs with the sysmon core of #2846), and their assertions only look at where the pass places Q/DQ nodes. The convs stay above the 16-channel limit of `graph_selection.py`, and both modes produce the same 8 Q and 8 DQ nodes as with the default size. `build_conv_concat_model` takes optional `channels`, `height` and `width` (defaults: the old values; with a seeded RNG the default model serializes to the same bytes). - **The autotune benchmark tests build a 1x32x32x32 model** instead of the 64x32x224x224 one (103 M random floats and about 1.2 GB copied per iteration; the five tests took 31 s, and none looks at the latency beyond "a positive float"). The two trtexec tests also pass `--duration=0`, because trtexec measures for at least 3 s whatever the iteration count. `_create_simple_conv_onnx_model` takes optional `batch_size` and `spatial_size` with the old defaults, so the unit autotuner tests that share it build the same model (their 161 tests pass). - **The autotune workflow test passes `--duration=0` to trtexec**: `test_export_quantized_model[True]` took 44 s, mostly inside trtexec's 3 s minimum measurement per scheme. It still runs the 100 requested iterations per measurement; the model and `timing_runs` keep their size, because the test asserts the best scheme carries Q/DQ nodes. ### Time improvement (measured in CI) gpu lane, time per file from the per-test timestamps; before = the runs of #2846-#2848, after = this PR's run and the 4 runs above it (median): | | before | after | change | |---|---|---|---| | three CPU-only files (68 cases), moved to the unit lane | 7.4 s | 0 | -7.4 s | | `test_concat_elim.py` | 15.2 s | 1.2 s | -14.0 s | | `autotune/test_benchmark.py` (36 cases) | 37.1 s | 21.0 s | -16.0 s | | `autotune/test_workflow.py` | 44.4 s | 28.5 s | -15.9 s | | **sum** | 104 s | 51 s | **-53 s per gpu job** (-59 s after speed scaling) | The whole `tests/gpu/onnx` directory takes 110-120 s before and 58 s in the three later runs on normal-speed runners (81 s and 74 s in #2849 and #2850, which ran on 1.3x slower runners). The 68 moved cases cost the unit job 7-13 s of CPU, spread over four xdist workers (below the unit job's run-to-run spread). My first estimate of the gpu saving (97-107 s) came from base runs without sysmon, where these files were slower. On its own, on main as of 07:06 UTC on 2026-10-21 (before #2846 added the sysmon core): the `gpu` lane's `tests/gpu/onnx` directory takes 145 s in #2854's run, which has the same base and lacks this change, and 66 s in this PR's run (-79 s). The runners were equally fast (median ratio 0.99 over the 17 untouched files above 20 s). `test_concat_elim.py` goes from 35 s to 2.7 s, `autotune/test_benchmark.py` from 39 s to 24 s and `autotune/test_workflow.py` from 46 s to 28 s, and the three moved files cost 14 s there. Times are the gaps between consecutive result lines in the job logs. 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 - The three moved files: 68 passed in 7.1 s at the old location and in 13.3 s from the new one (this includes the unit conftest's session setup) on Python 3.12 with onnx 1.21.0 and onnxruntime 1.24.4, and 68 passed in 15.3 s on Python 3.10.19 with onnxruntime 1.22.0, the oldest runtime `pyproject` allows. The whole `tests/unit/onnx` directory passes (756 passed, 12 xfailed). - `test_concat_elim`: the two tests pass in 1.3 s on CPU; quantizing the default and the small model in both modes gives Q=8, DQ=8. - The shared helper changes: the unit autotune tests (161) pass. - Not run: the TensorRT tests, because there is no TensorRT here; this PR's `gpu` job is the check. Pre-commit 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 - Did you write any new necessary tests?: N/A (existing tests moved or made smaller; the 68 moved cases are unchanged) - 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 5 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. #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist (merged) 2. #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI (merged) 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 (targets main) **(this PR)** 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 (merged) 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 (merged) 11. #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Tests** * Added configurable dimensions to ONNX test models and smaller inputs for selected scenarios. * Streamlined GPU benchmark timing while retaining configured measurement runs. * Updated INT8 and FP8 concat-elimination coverage to use CPU calibration and smaller model shapes. * Reorganized DLA and quantization checks as unit tests instead of GPU tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: CI and test runtime (reference draft for runtime measurement, not for merge)
This branch holds the whole GPU-CI speed-up series (34 commits, 193 added lines of non-test source, no
modelopt/library code touched), rebased on main786b84e4a5, so that real runner runs could measure it. It is now split into 11 stacked draft PRs (stack #2859), each a small functional change that stands on its own. Three changes of this draft that saved nothing measurable (thecoverage.xml-only report,mmlu_limit=1fortest_qwen3_eval_fp8, the vLLM KV-pool sizing) are not in the stack; this branch still has them. A separate PR based onmain, #2858, makes the required checks run in GitHub's merge queue.Where the time goes today (successful PR runs, 2026-09-25 to 2026-10-09): the median GPU run of a PR takes 59 min, of which 13.8 min is waiting for the unit
linuxjob and 33.5 min is pytest; gpu_megatron starts its worker pool 15 times (about 76 s each); thegpulane compiles three CUDA packages in every job (about 7 min); the nightly 2-GPUgpujob hit its 60-min cap on 15 of 15 nights.What the stack leaves out, on purpose. #2836 (open) also caches the source-built wheels of the
gpujob, saves the JIT cache when a job fails, and destroys every non-default NCCL group in the Megatron teardown; its wheelhouse is the more complete one (pinned wheels, C++20 patch, saved from any run). So this draft's wheel-cache and JIT-save commits are not in the stack, and neither are the Triton cache and per-lane cache keys, which build on #2836's save action and can follow once it lands. The group cleanup in #2855 is the same loop as #2836's, so whichever lands second sees an identical hunk.ExampleRunner(a worker runner for the torch_onnx lane, +1,100 lines with its tests) is a third in-process example runner next to Megatron's and the diffusers one (#2728): it is not in the stack until #2728's runner has been compared on torch_onnx. The numbers below therefore include pieces the stack does not carry (the wheel and Triton caches, the runner).Stack results (2026-10-17). Every PR of the stack was green on all required checks (Unit, Code Quality, GPU, Example, Regression, DCO and the four aggregate checks; Docs where it runs) on its previous head. Dropping the three changes rewrote the heads of #2846-#2855. All of them ran again on the new heads and are green (Unit, Code Quality, GPU, Example, Regression, DCO); the table below and the measurements are from the earlier runs, whose only difference to the new heads is those three changes.
build-docs, which the ruleset also requires, is reported on #2845 but not on #2846-#2855: the Docs workflow did not trigger when their branches were force-pushed and re-stacked.codecov/projectfails on #2851 only because that PR changes example tests alone, so path gating skips the GPU lanes and Codecov compares a partial upload against a full base; it is not a required check. Minutes of each PR's own CI run (one run per PR,-= lane skipped by path gating); the stack is cumulative, so #2855 contains all the others:linuxjobgpujobgpu_megatronjobRead the table above as a trend, not as per-PR attribution: single runs are noisy (the
gpulane swings between 30 and 47 min with whether the JIT caches hit, which depends on the PR branch and the container image), and the Example lanes include other contributors' changes on main since the baseline (hf_ptq, diffusers). The per-PR attribution is measured in the next table.Measured improvement per PR (from the stack's own CI runs, with the PR below as the control, speed-scaled medians; each PR's description has its table, the control runs and the method; one test varies by 15-30% from run to run):
linuxjob 12.9 min median serial (27 runs of other PRs) to 7.0 min median over the stack's runs (with sysmon from #2846; this PR's own run: 10.9 min); unit workflow 24.1 to 17.2 minThe first version of the PR descriptions estimated more than the runs show in some places and less in others. Three changes saved nothing measurable and were dropped from the stack on 2026-10-17: the
coverage.xml-only report (no CI effect, the pytest summary line comes 1-4 s after the last test either way),mmlu_limit=1fortest_qwen3_eval_fp8(187 to 176 s, noise) and the vLLM KV-pool sizing (about 5 s per vLLM version, noise). The sysmon core of #2846 had also made the ONNX tests of #2849 cheaper than first estimated (53 s instead of 97-107 s) while the local-data tests of #2850 saved more (170 s instead of 94 s). The descriptions now carry the measured numbers. The second set of runs (the rebuilt heads) is a replicate of the same code. It reproduces the direction of every effect, and two of them to a percent: the torch_onnx lane (-48% in the first runs, -49% in the second) and the Megatron session pool (-21%, -22%). The example-lane trims of #2851 and the DFlash lane of #2852 come out smaller in the second set (llm_qat -20% and -19%, megatron_bridge -21% and -24%, llm_eval -36% and -18%, speculative_decoding -16% and -8%, DFlash -38% and -24%): the later runs landed on slower hosts, and the same lane varies by up to 50% between the two sets, so the lane numbers carry about +-20-30%. The merge queue PR #2858 is separate and not part of these measurements: a queued merge takes about as long as the GPU workflow (59 min today, 36-40 min with the series).The first CI runs found one real problem: #2855 (the session-wide pool) failed in gpu_megatron twice, because a failed job's poisoned CUDA context was reused and because its teardown did not collect garbage after destroying the groups; both are fixed in that PR (see its description).
Usage
Nothing user-facing. The series is CI and test changes apart from
examples/torch_onnx/torch_quant_to_onnx.py, whose calibration loaders now usenum_workers=0and whosebuild_trt_engineruns one benchmark iteration.Testing
All six workflows are green on this head (Unit with 20 jobs, Docs, Code Quality, GPU, Example, Regression). Run locally, CPU only: pre-commit (every hook, including mypy) and actionlint over the range;
tests/unit(withoutpuzzletron, which needshydra) in normal, reversed and shuffled file order, 5,047 passed and one failure in a file this series does not touch (fastgen/test_qwen_image_pdd_plugin.py::test_mr210_joint_mask_ignores_padded_token_values, local environment only), on the tip before the last rebase (same patches, one upstream commit older); collection oftests/unitwith the POSIX-onlysignal/osnames removed (Windows emulation).Minutes. The first two runs are on the previous head of this branch (
524d784d71, baseaf07754681); the third is this head on main786b84e4a5, where thegpulane moved to a new container image (#2809) and its JIT caches missed, and the other lanes' caches hit only in part:linuxjobThe
gpulane did not improve on this head: it ran with missed JIT caches and still compiles the three CUDA packages, since the wheel cache can only be seeded frommain(what #2836 addresses). The example lanes (hf_ptq, diffusers, llm_eval, torch_onnx and others) also include other contributors' changes on main since the baseline.Before your PR is "Ready for review"
Draft, not for merge. Commits are signed (
git commit -s -S).CONTRIBUTING.md: no new dependency (pytest-xdistis already in thedev-testextra on main)tests/gpu/torch/utils/test_distributed_worker_pool.pyAdditional Information
Reference only. The reviewable PRs are the stack above.