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>
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
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>
shengliangxu
added a commit
that referenced
this pull request
Oct 21, 2026
### What does this PR do?
Type of change: CI (test runtime); the only non-test file is
`examples/torch_onnx/torch_quant_to_onnx.py`
Three cuts to the torch_onnx example lane:
- **No DataLoader workers in the calibration loaders.**
`load_calibration_data` builds its DataLoaders over lists of tensors
that are already on the GPU, so the four workers load nothing, and since
the script forced the `spawn` start method (CUDA cannot be used in
forked workers), every loader pass spawned four processes that
re-imported torch, timm and modelopt and received the dataset over CUDA
IPC: about 14 s per pass and about 300 s per torch_onnx CI job. Both
loaders use `num_workers=0`, and the module-level
`mp.set_start_method("spawn", force=True)` is dropped, which also stops
importing the script from changing the start method of the importing
process (two tests import it inside the pytest process). Nothing else in
the script uses multiprocessing. The `--evaluate` path keeps the
CPU-only DataLoader workers of `examples/onnx_ptq/evaluation.py`, now
under the default start method, exactly as in the `onnx_ptq` example
that calls the same function.
- **swinv2 int8, mxfp8 and auto are skipped on pull requests.**
`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/`) fp8 and nvfp4 still run
and 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 three skipped cases cost about 110
s of every PR run.
- **`build_trt_engine` runs trtexec for one iteration.** It let trtexec
run its default benchmark after the build (200 ms warm-up, then at least
3 s and 10 iterations) and discarded the output. `--warmUp=0
--duration=0 --iterations=1` still executes the engine once and removes
about 3 s per engine, roughly 40 s per lane run. `--trt_build` users see
no difference besides it returning sooner.
### Time improvement (measured in CI)
torch_onnx lane; before = #2846, #2848 and #2850 (875-880 s; #2849 ran
on a slower runner: 1,060 s), after = #2853, #2854 and #2855 (452-463
s).
| | before | after | change |
|---|---|---|---|
| **lane pytest** | 877 s | 458 s | **-419 s (-48%)** |
| of which: the three swinv2 cases skipped on PRs | 143 s | 0 | -143 s |
| of which: the cases that still run (no DataLoader workers, trtexec
runs once) | | | -277 s |
| whole job | 18.6 min | 12.2 min | -6.4 min |
This lane repeats to within 1% (875-880 s before, 452-463 s after), so
it is the cleanest measurement of the series. For example `swin_tiny`
with auto quantization goes from 66 s to 24 s and `vit_tiny` with auto
from 60 s to 17 s.
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
No API change. The skip is keyed on `GITHUB_REF`, so a local run keeps
every combination.
### Testing
`py_compile`, ruff 0.15.20 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 (it needs a
GPU, TensorRT and model downloads) and the flags against a real trtexec,
so this PR's torch_onnx job is the check. Pre-commit passes. #2685
(open) also edits this example; the changes touch different lines.
### 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; the
PR-only skip keeps test ids)
- 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 9 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, 5, 8 and 10 (#2845, #2846, #2849, #2852,
#2854) have merged. The remaining slices (#2847, #2848, #2850, #2851,
#2853, #2855) are independent PRs on main: each applies to main on its
own, and merging all six onto main in different orders gives the same
tree. 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 (merged)
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 **(this
PR)**
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
-->
## Summary by CodeRabbit
* **Updates**
* Quantization examples now use zero-worker calibration loading with
device-resident samples.
* TensorRT example builds perform a single inference iteration without
warmup or timing.
* **Tests**
* Pull-request validation skips selected SwinV2 Tiny formats, while
scheduled validation continues to cover them.
* Pull-request detection uses the repository’s CI reference environment.
* **Impact**
* These changes affect example workflows and validation coverage only;
no application user-facing functionality is 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
… tests, shrink puzzletron data (#2850) ### What does this PR do? Type of change: CI (test runtime and reliability, tests only) Less data and fewer Hub round trips in four GPU-lane tests: - **Six `preprocess_data` tests use the tiny local tokenizer** (`tests/gpu_megatron`). They tokenized with `gpt2` from the Hub, and `megatron_preprocess_data` loads the tokenizer once for the vocabulary size and again in every worker, so each test made several round trips; none depends on gpt2's vocabulary. They use the checked-in byte-level `tiny_tokenizer_path` fixture that the blend test already uses. Three of the six read only local data, so the tokenizer download was their only network access (the jsonl `max_tokens` test is among the lane's 120 s timeouts, 3 failures). Measured on a CPU host with a cold Hub cache, the 11 tests of the module take 42.8 and 44.0 s before and 35.0 and 36.1 s after (two runs each; the nine tests of the earlier base took 38.5 s and 24.4 s, and 37 s and 23 s in CI). The Qwen3 tests keep their real tokenizer (chat template and EOS handling). - **The puzzletron test data shrinks, and a duplicate NAS test is folded into another.** `save_dummy_dataset` wrote 2500 copies of one conversation into each split (1738 blocks of 2048 tokens), while the pruning scores read 100 blocks and the later steps 2 each; 400 conversations (278 blocks) take the three dataloader builds plus the save from 2.8 s to 0.36 s per test, and the validation loss moved by at most 0.02 on random-init toy Llamas (the scoring tolerance is 0.15). `calc_subblock_stats` ran 12 passes where only the entry for `batch_size` 96 is read: `batch_sizes: [96]` runs 4. `test_nas_convert_ffn_pruning` ran the same `mtn.convert` on the same toy Llama and config as `test_nas_search` (21.6 s with its own 2-rank spawn); its two assertions move into `test_nas_search`. - **The unified HF export test calibrates on local text, with no previews or batch probing.** Each of its 11 params starts a fresh `hf_ptq` process (22 s per param). It now reads a local jsonl of 64 random-word texts (new helper `write_calib_jsonl`, shared through the session-wide `calib_jsonl` fixture in `tests/gpu/conftest.py`, which fails the session if a test writes into its directory) instead of `cnn_dailymail` from the Hub, passes `--batch_size 64` (the whole calibration set in one batch instead of probing for the largest batch) and `--skip_generate`. The test only checks the structure of the exported safetensors. - **The GPTQ end-to-end test reads 8 texts from the same local jsonl** instead of streaming `cnn_dailymail` for each of its three params (and a test that could fail on a Hub hiccup). ### Time improvement (measured in CI) Time of the changed tests; before = the runs of #2846-#2849, after = #2850, #2852, #2854 and #2855 (median). #2849 and #2850 ran on runners 1.26x and 1.29x slower for every test, so the change after speed scaling is in brackets. | lane | tests | before | after | change | |---|---|---|---|---| | gpu | `tests/gpu/torch/puzzletron` (18 tests, 17 after the folded NAS test) | 290 s | 202 s | -88 s (-70 s, -26%) | | gpu | unified HF export, 11 params | 228 s | 153 s | -75 s (-67 s, -30%) | | gpu | GPTQ end-to-end, 3 params | 7.2 s | 1.1 s | -6 s | | gpu_megatron | `preprocess_data` module | 37 s | 23 s | -14 s (-10 s, -27%) | So about 170 s (144 s after scaling) off the gpu job and 14 s off the gpu_megatron job. This PR's own run was on a slow runner (puzzletron 310 s); the three later runs agree with each other (199-203 s). My first estimates were 50 s, 38 s and 6 s for the first three rows. 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 - `preprocess_data`: ran the module on CPU next to a copy of the `tiny_tokenizer_path` fixture with the Hub reachable: 11 passed before and after (main has dropped the HF xfail marks), with identical test ids. - `calib_jsonl`: with the real fixture and the repo's `get_dataset_dataloader` on a CPU host, the 8 samples read from the shared 64-text file equal those of the old 8-text file, and two reads leave the directory unmodified (the fixture's teardown check passes). `--setup-plan` shows the unified export test using the conftest fixture; the GPTQ and export tests themselves need a GPU. - puzzletron: timed the dataloader builds, the subblock-stats passes and the loss sensitivity on CPU with the repo's own helpers; collected `tests/gpu/torch/puzzletron` (16 tests, one fewer than before). 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. - unified export: no GPU here, so I ran the test on a CPU host through a 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 need CUDA or Triton. GPTQ: the flow completes on the jsonl for the FP8 and INT4 blockwise configs (NVFP4 needs CUDA). - Not run: the GPU tests themselves; this PR's gpu and gpu_megatron jobs are 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?: ✅ one small helper, `write_calib_jsonl` (`tests/_test_utils/torch/calib_data.py`), and one session fixture, `calib_jsonl` (`tests/gpu/conftest.py`), used by two tests; no new test functions - 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 6 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, 5, 8, 9 and 10 (#2845, #2846, #2849, #2852, #2853, #2854) have merged. The remaining slices (#2847, #2848, #2850, #2851, #2855) are independent PRs on main: each applies to main on its own, and merging all of them onto main in different orders gives the same tree. 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 (merged) 6. #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data **(this PR)** 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 (merged) 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** * Updated model export and quantization checks to use locally generated calibration data. * Standardized preprocessing tests on a small local tokenizer. * Restricted model analysis configurations to a batch size of 96. * Reduced sample dataset sizes for pruning-related checks. * Strengthened conversion verification by checking generated pruning outputs. * Removed the FFN pruning conversion test while retaining attention-pruning coverage. <!-- 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
…ridge example tests (#2851) ### What does this PR do? Type of change: CI (test runtime, tests only) Five cuts to the example lanes, each with its reasoning: - **`test_lm_eval_hf` evaluates one MMLU subtask.** It ran `--tasks mmlu --limit 0.1`: the 57-subtask group, about 5.6k loglikelihood requests, on a tiny Qwen3 whose `max_position_embeddings` is 32, so its runtime is dominated by loading one HF dataset config per subtask (over 331 CI runs: min 198.9 s, median 212.5 s). It runs `mmlu_abstract_algebra` (100 test documents, 5-shot) through the same `lm_eval_hf.py` path and the same accuracy gate (exactly one task, lower bound 0.1 unchanged; the label distribution of that split is 22/26/31/21, so even a model that always predicts one letter scores at least 0.21). - **The llm_qat QAT and QAD tests quantize once per module.** `test_qwen3_qat_nvfp4` (3 backends), `test_qwen3_qad_nvfp4` (2) and `test_qwen3_qlora_nvfp4` ran 6 `quantize.py` and 12 cold tokenized-dataset builds per module (104-121 s each, qlora 169 s). A module-scoped fixture runs `quantize.py` once (guarded by `assert_unmodified_tree`, like the other shared tiny-model fixtures) and all quantize and train calls share one dataset cache; the qlora test keeps its own PTQ because it passes `--compress True`. - **The 1-GPU EAGLE3 compile smoke test skips the Triton GEMM autotune** (`TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=ATEN`, the mechanism `test_diffusers.py` already uses). `test_llama_eagle3[1-False]` takes 175.5 s against 40.6 s for `[1-True]` (the same flow without compile), about 135 s of it `torch.compile` benchmarking Triton GEMM templates for a draft model that is one tiny decoder layer. The mode stays `max-autotune`, and the shortcut applies only when the run has one GPU: `[2-False]` and any run on more GPUs (the nightly) keep the full autotune. - **The duplicate tiny-llama offline EAGLE3 training test is dropped** (`test_offline_eagle3_training[tinyllama]`, 34.9 s per job, almost all launcher start-up). `test_offline_eagle_training` already runs that mode on the same model with hidden states from `compute_hidden_states_hf.py`, and the FakeBaseModel case keeps the synthetic path. The one thing the dropped case covered that stage 2 did not, the collator truncating dumps, moves there (`training_seq_len=16`, so 28-token dumps are truncated), and the collator is also unit tested on CPU. - **`megatron_bridge` tests that use no GPU skip the Megatron state reset.** The conftest resets Megatron's global state around every test (`destroy_model_parallel`, `gc.collect()`, `empty_cache`), about 1.3 s per test; 77 tests in the directory are pure Python (the mlflow wiring and checkpoint-path tests) and leave nothing to reset (the 82 tests that took under 3 s each add up to 118 s in the series' runs before this PR and to 0.3 s after it). The reset now runs only for tests that request `num_gpus`; the environment restore stays for every test. The job runs at 68-75 minutes against its 75-minute nightly limit. ### Time improvement (measured in CI) Example lanes, pytest time of the changed tests and of the lane; before = #2846 (plus #2848 where its runner was as fast), after = #2851, #2854 and #2855 (median). #2849 and #2850 ran on runners 1.3-1.5x slower for every test and are left out of the before column. | lane | | before | after | change | |---|---|---|---|---| | llm_eval | `test_lm_eval_hf` | 130 s | 19 s | -111 s | | llm_eval | **lane** | 332 s | 212 s | **-120 s (-36%)** | | llm_qat | six QAT, QAD and QLoRA tests (call time) | 563 s | 428 s | -135 s (the shared quantize adds 20 s to the first test's setup) | | llm_qat | **lane** | 591 s | 475 s | **-116 s (-20%)** | | speculative_decoding | `test_llama_eagle3[1-False]` | 131 s | 87 s | -44 s | | speculative_decoding | dropped `test_offline_eagle3_training[tinyllama]` | 20 s | - | -20 s | | speculative_decoding | **lane** | 407 s | 341 s | **-66 s (-16%)** | | megatron | the 82 megatron_bridge tests that use no GPU (the state reset) | 118 s | 0.3 s | -118 s | | megatron | **lane** | 583 s | 460 s | **-123 s (-21%)** | - A `mmlu_limit=1` change for `test_qwen3_eval_fp8` was dropped from this PR. The runs above are of its earlier head, which had it: the test took 175-246 s in all 7 runs, before and after (187 s against 176 s on comparable runners), so the CI history it was based on (mean 723.7 s, 2026-09-25 to 10-09) is not reproduced by the current lane and the change saved nothing measurable. The llm_eval lane figure above is `test_lm_eval_hf`'s. - That is 425 s of lane time per PR run. The Example workflow ends with its slowest lane, torch_onnx (18.1-18.8 min of job time) until #2853, so the gain reaches the workflow's wall time only with that PR. 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 - llm_eval: ran the test's command on CPU with the tiny Qwen3 it builds: old configuration 223.4 s, `mmlu/acc` 0.2343; new 22.4-24.0 s, `mmlu_abstract_algebra/acc` 0.27, gate passes. - llm_qat: with `run_example_command` faked, `quantize.py` runs 6 -> 2 times and cold dataset builds go 12 -> 2; real `quantize.py` and `train.py` (QAT, then QAD from the checkpoint) ran on CPU with the int4 and fp8 recipes (NVFP4 needs CUDA) and left the checkpoint unmodified. - EAGLE: stage 1 and 2 of the offline flow ran on CPU with `training_seq_len` 64 and 16 (the collator pads 28 -> 64 and truncates 28 -> 16); with `run_example_command` faked, only `[1-False]` on one GPU gets the autotune variable (that case on 2 GPUs and every other case keep the full autotune) and the only removed test id is `test_offline_eagle3_training[tinyllama]`. - megatron_bridge: a throwaway pytest module (not in this PR) against the patched conftest counts the resets: plain tests trigger none, tests that request `num_gpus` trigger one before and one after; it fails against the old conftest. - Not run: the example lanes themselves (they need GPUs, TRT-LLM and model downloads), so the accuracy with CUDA bf16 numerics, the savings and the autotune variable are checked by this PR's example jobs. Pre-commit passes. - `trust_remote_code=True` in `test_offline_eagle3_training` is not new: it is the line on main (the Kimi and MiniMax repos it loads need remote code), and this PR only renames its variable after dropping the tinyllama case. ### 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; one duplicate test dropped, see the description) - 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 7 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, 5, 8, 9 and 10 (#2845, #2846, #2849, #2852, #2853, #2854) have merged. The remaining slices (#2847, #2848, #2850, #2851, #2855) are independent PRs on main: each applies to main on its own, and merging all of them onto main in different orders gives the same tree. 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 (merged) 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 **(this PR)** 8. #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test (merged) 9. #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane (merged) 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** * Focused language-model evaluation coverage on the abstract algebra task while retaining five-shot evaluation. * Streamlined quantization-aware training scenarios by reusing cached datasets and checkpoints, reducing repeated setup. * Improved Megatron test isolation by applying state cleanup only where needed. * Refined EAGLE test configurations for single-GPU runs, remote model coverage, and offline training with shorter sequence lengths. <!-- 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
…2848) ### What does this PR do? Type of change: CI (test reliability) `DistributedWorkerPool` waited 120 s for its workers to import the Megatron/TE/modelopt stack and report ready. The start takes 75-77 s in 14 modules of the PR run I read and has a p90 of 94-98 s, so slow hosts time out and the whole module errors: 15 jobs (903 job-minutes) in two weeks, 15 of the 63 failed gpu_megatron PR jobs in that period. It waits 300 s. An earlier version of this PR also gave `test_mamba_search_space` more time. Main's #2836 has since set its cap to 150 s, and that holds: in 58 gpu_megatron jobs from the last six days the test took 64-132 s (median 73 s, none above 150 s), so nothing is left to change there. ### Time improvement (measured in CI) None in a passing run: this PR raises one limit that only failing runs reach. - In the 9 gpu_megatron runs of this series, the worker-pool start took 75-77 s before #2846 and 18-27 s after it, against the old 120 s ready timeout. No run came near the limit, so the series cannot show the failures this PR prevents (the sysmon core of #2846 also takes most starts out of reach of the old ready timeout). - What it prevents, from the CI history of the two weeks before the series: the ready timeout failed 15 jobs (903 job-minutes), 15 of the 63 failed gpu_megatron PR jobs. Each failed job costs the author a re-run of the lane, 20-45 min of job time in this series. Method: every PR of the series has its own CI run, so "before" is the run of the PR(s) below that lack the change and "after" the runs that carry it. Times come from the job logs (`--durations`, or the gap between consecutive result lines). Where runner speed matters, the tests that no PR of the series touches scale each run, and the raw and the scaled medians are both shown. One test varies by 15-30% from run to run on its own, so smaller differences are not evidence. ### Usage Nothing user-facing. ### Testing Pre-commit passes. Not run on a GPU here: the effect is on cold-cache runs, so this PR's gpu_megatron job can only show that nothing regresses. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](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 (a timeout in the shared test utilities) - 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 4 of 11 of the GPU-CI speed-up work, split out of the reference draft #2730 (not for merge; it holds the whole series with its measurements). Slices 1, 2, 5, 8, 9 and 10 (#2845, #2846, #2849, #2852, #2853, #2854) have merged. The remaining slices (#2847, #2848, #2850, #2851, #2855) are independent PRs on main: each applies to main on its own, and merging all of them onto main in different orders gives the same tree. 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 **(this PR)** 5. #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models (merged) 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 (merged) 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 * **Bug Fixes** * Increased the wait time for distributed worker initialization to improve reliability during startup. <!-- 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
…egatron modules (#2855) ### What does this PR do? Type of change: CI (test runtime); changes `tests/gpu_megatron` and `tests/_test_utils` only Each of the 15 gpu_megatron modules that use `dist_workers` started its own worker pool: about 19 of the lane's 36 minutes. Fourteen logged module starts took 75-77 s each (the GDN module 179 s: 76 s pool plus 103 s Triton JIT), plus 3.6-10 s of shutdown, and the start has a p90 of 94-98 s. The 76 s is not distributed init (an in-process init takes 0.1 s): the spawned child unpickles `megatron_worker_teardown` by reference, which imports the heavy `tests/gpu_megatron/conftest.py` (transformers, megatron.core, all of modelopt and its plugins, TE) before it reaches `init_process_group`. The same import chain takes pytest itself 89 s to start here against 17 s in the gpu lane, whose session-scoped pool, with a light teardown function, starts in 1.7 s. Module scope was chosen when a spawn cost about 10 s ("avoids cross-test state contamination"); the git history shows no contamination incident, and the cost has grown 7.6x. `_pool_cache` is now session-scoped, so the pools are keyed by world size and shared by every module. The `dist_workers` fixtures stay function-scoped and get their pool from `get_live_pool`, which replaces a pool that a hang killed. What isolation the module scope gave is kept explicitly: - `megatron_worker_teardown`, after every job: synchronizes CUDA first (an asynchronous CUDA error then surfaces in the test that caused it), destroys Megatron's model-parallel state as before, and now also the process groups it leaves behind (about 20 NCCL groups per dispatch, each with a communicator and a watchdog thread; `test_megatron.py` alone dispatches about 110 times). A pool that lives through the whole session needs it more than a per-module pool did. - `_reset_shared_pools`, an autouse module-scoped fixture: when a module starts, it resets the process-global torch settings that helpers like `set_seed` change (cudnn and TF32 flags, deterministic algorithms, RNG) to what a fresh worker starts with (`reset_worker_state`; the workers record those settings right after they start), in the pools that already exist. It is not part of `_get_pool`, which runs once per test and would add a job to each. - `get_live_pool` and the pool's abort handling: a job that times out, is interrupted, loses a worker **or fails on any rank** retires the pool, and the next test gets a fresh one, for the cost of one pool start. `DistributedWorkerPool.run` has retired a pool after a plain failed job since #2836, which matters more with a session-wide pool: a CUDA error such as the known Blackwell illegal memory access is sticky, the workers stay poisoned, and every later test would fail with it. `_reset_shared_pools` no longer lets a failed reset error the module's tests (they may not use a pool at all). Measured saving on top of #2846's sysmon core: 216 s of the gpu_megatron lane (1,038 s to 822 s; see the time section). Without sysmon a pool start costs 76 s, and the saving would be about 17 min (an estimate; the series was not run that way). One pool start instead of 15 also means fewer chances to hit the ready timeout. ### Time improvement (measured in CI) gpu_megatron lane, pytest time: | main (#2845) | + #2846 | slices 1-10 (the earlier base) | + this PR | |---|---|---|---| | 2,240 s | 1,212 s | 1,038 s (1,029 s in #2852) | **822 s** | This PR on its own: 1,038 s -> 822 s, **-216 s (-21%)**, against the two earlier runs (#2852: 1,029 s, #2854: 1,038 s). This PR's run was on a slightly slower runner (the heavy `call` phases of the tests it does not change took 6% longer than in #2854), so the gain is not runner luck. The setups of 12 s or more, which are the pool starts and the module-level compiles, fall from 16 (489 s) to 3 (246 s): the pool starts once instead of in each of 14 modules, and what remains is the extension build (124 s), the GDN module's Triton JIT (105 s) and one 18 s pool start. With sysmon (#2846) a pool start takes 18-27 s instead of 75-77 s, so each avoided start saves less than the 76 s this PR was designed around. Without #2846 the saving would be 14 starts of 76 s plus the 3.6-10 s shutdowns, about 17 min; that is an estimate, because the series was never run without sysmon on this PR. This is one CI run of the final head: the two runs before it failed on the sticky CUDA error and on the missing `gc.collect()` described below. Method: every PR of the series has its own CI run, so "before" is the run of the PR(s) below that lack the change and "after" the runs that carry it. Times come from the job logs (`--durations`, or the gap between consecutive result lines). Where runner speed matters, the tests that no PR of the series touches scale each run, and the raw and the scaled medians are both shown. One test varies by 15-30% from run to run on its own, so smaller differences are not evidence. ### Usage Nothing user-facing (test fixtures). ### Testing - `reset_worker_state` has a gloo test next to the other pool tests (it runs on CPU): flags and seed changed in a job are back at their spawn-time values afterwards; it fails when the helper forgets one of the settings. - The Megatron lane needs GPUs, so this PR's gpu_megatron job is the check: 822 s of pytest on the final head, against 1,038 s for the earlier base, #2854's run (381 passed, 129 skipped, as before). - The first CI run of this PR failed in gpu_megatron: the Blackwell illegal-memory-access flake hit `test_lora_then_quantize` and poisoned the shared pool, and every later module errored (3 failed, 180 errors), because the pool was reused after a failed job and `_reset_shared_pools` ran a job in it at each module start. The second commit retired a pool after any failed job and made a failed reset harmless; with it the lane ended with 1 failure instead of 183, in `test_forward_with_one_lora[lora_config1]` (the same CUDA error). Re-running the job on that commit failed at the same test again, so it was not a random flake: the teardown of this PR, unlike the draft's, did not collect garbage after destroying the groups, and the draft passed the lane three times. The third commit adds that `gc.collect()`; the lane then passes. The cause is not established (the finished test's Megatron and TE modules can still reference the destroyed groups, and their finalizers then run at an arbitrary point of the next job). #2836 has since landed the retirement in `run()` for its own reasons (with its own gloo test), so the second commit now only makes a failed reset harmless; main's teardown has neither the process-group loop nor the collection, and a session-wide pool needs both. - Pre-commit passes. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](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?: ✅ 1 test function in `tests/gpu/torch/utils/test_distributed_worker_pool.py` (gloo, CPU) for the new helper; the session pool itself is exercised by the whole gpu_megatron lane - Did you update [Changelog](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 11 of 11 of the GPU-CI speed-up work, split out of the reference draft #2730 (not for merge; it holds the whole series with its measurements). Slices 1, 2, 4, 5, 6, 7, 8, 9 and 10 (#2845, #2846, #2848, #2849, #2850, #2851, #2852, #2853, #2854) have merged. The two remaining slices (#2847, #2855) are independent PRs on main: each applies to main on its own. 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 (merged) 5. #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models (merged) 6. #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data (merged) 7. #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests (merged) 8. #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test (merged) 9. #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane (merged) 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 PR)** <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Tests** * Distributed GPU test workers now restore their initial PyTorch settings and random seed between test modules. * Worker teardown now synchronizes CUDA, removes non-default process groups, and cleans up Python and GPU resources. * Worker pools are shared across test modules and shut down at the end of the test session for more consistent test execution. <!-- 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
… GPU runs (#2847) ### What does this PR do? Type of change: CI (test runtime) Tests marked `manual` collect only when pytest gets `--run-manual`. The unit workflow passes it for every run that is not a pull request, but the GPU workflow never did, so a `manual` GPU test never ran in CI. This sets `PYTEST_ADDOPTS=--run-manual` in the "Run gpu tests" step for every run that is not a pull-request run (the nightly schedule and manual dispatches), which lets slow or network-heavy checks leave the pull-request jobs without leaving CI. Pull-request runs are unchanged by the flag itself: copy-pr-bot runs them as `push` events, the workflow's only `push` trigger, and the expression is `github.event_name != 'push' && '--run-manual' || ''`. An earlier version of this PR used the `release` marker; review feedback pointed out that `release` is the QA team's marker for long real-model runs in their own pipeline, while `manual` is what the unit tests already use for nightly and on-demand runs. **Side effect:** `--run-manual` also runs the existing `manual`-marked tests of `tests/gpu` on those runs, which no GPU run executed before: 12 NAS vision-model cases (`test_search_space_with_vision_models.py`), 6 puzzletron cases (about 40 s each by their marker) and the TRT client benchmark models (they need TensorRT, so I could not count them on a CPU host). The unit lane already runs its `manual` tests nightly the same way. A `workflow_dispatch` run of the GPU workflow on this branch (run 39668011982: 2-GPU runners and `--run-manual`, the nightly's setup) passed all four GPU jobs: `gpu` in 28 min (1,866 passed, 40 skipped; cap 60 min) and `gpu_megatron` in 27 min (403 passed, 87 skipped; cap 75 min). In `gpu` the six puzzletron params took 174 s (26-32 s each), the 14 NAS vision results 29 s in total and the Nemotron v3 test 20 s, about 4 min together; the TRT client tests are skipped in that container (`libnvinfer.so.10` is missing), so they add nothing. Two tests use it: - `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 21-39 s call in every PR gpu_megatron job (21 s with the sysmon core of #2846, 39 s without it), and the only test in the lane that downloads real weights. - `test_get_dataset_dataloader_nemotron_v3_chat_template` streams the seven gated Nemotron post-training v3 datasets from the Hub (a 17-24 s call that needs an HF token and the network, and a source of Hub flakes). The combo expansion and sample distribution it goes through are covered by CPU unit tests that record the per-dataset calls. **Trade-off:** these two checks no longer run on pull requests. The real-weight accuracy check (0.36 < MMLU < 0.39) and the gated-data smoke test run every night and on demand, a day after a change instead of on the PR. ### Time improvement (measured in CI) Time of the module that holds the skipped test, from the per-test timestamps; before = the run of #2846, after = this PR's run and the runs above it that carry the change (median): | lane | skipped test (call time before) | module before | module after | change | |---|---|---|---|---| | gpu_megatron | `test_megatron_generate_and_mmlu[tp]` (21 s) | 64 s | 41 s (39 s in this PR's run) | -22 s (-35% after speed scaling) | | gpu | `test_get_dataset_dataloader_nemotron_v3_chat_template` (17 s) | 32 s | 13.5 s (13 s in this PR's run) | -19 s (-62%) | So 22 s off the gpu_megatron job and 19 s off the gpu job of every PR run. The lane totals cannot show a change this small (gpu_megatron 1,212 -> 1,206 s, gpu 1,289 -> 1,351 s: their run-to-run spread is several times larger). The price is that the two checks run only on the nightly and on demand. 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 `pytest --run-manual` runs the manual-marked tests locally. ### Testing On a CPU host the Nemotron v3 test is skipped with "Only run when --run-manual is given" without the flag and runs with it, and `-m manual` selects it. The expression `github.event_name != 'push'` is false for the `push` events that run pull requests and true for the nightly and manual dispatches. Pre-commit and actionlint pass. Run on GPU runners with the flag (dispatch run 39668011982, the nightly's setup): the Nemotron v3 test passed in 20 s and the four Megatron MMLU cases in 145 s (tp 66 s, cp 39 s, pp 32 s, dp 8 s). This PR's own GPU run shows both skipped on a pull request. ### 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 (marks two existing tests; no new tests) - 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 3 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, 4, 5, 6, 7, 8, 9 and 10 (#2845, #2846, #2848, #2849, #2850, #2851, #2852, #2853, #2854) have merged. The two remaining slices (#2847, #2855) are independent PRs on main: each applies to main on its own. 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 **(this PR)** 4. #2848 [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time (merged) 5. #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models (merged) 6. #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data (merged) 7. #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests (merged) 8. #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test (merged) 9. #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane (merged) 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** - GPU test coverage now includes selected manual tests during scheduled and on-demand workflow runs. - Tests requiring gated datasets or model downloads are excluded from standard push-based runs to keep routine checks focused and efficient. - **Chores** - Updated GPU test execution settings to support explicitly enabled manual test runs. <!-- 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.