Repository navigation
[Draft, not for merge] GPU CI speed-up series (runtime measurement) - #2730
Draft
shengliangxu wants to merge 34 commits into
Draft
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>
test_lm_eval_hf 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 every context is truncated to 32 tokens. Its runtime is dominated by loading one HF dataset config per subtask and does not depend on the host (over 331 CI runs: min 198.9 s, median 212.5 s, mean 221.0 s). Run `mmlu_abstract_algebra` (100 test documents, 5-shot) instead. It goes through the same lm_eval_hf.py code path and the same accuracy gate (which requires exactly one task and reads results[task]['acc,none']); only the number of datasets and documents changes. The gate's lower bound of 0.1 is 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. Testing: ran the test's command on CPU (CUDA hidden) with the tiny Qwen3 the test builds. Old configuration: 223.4 s wall, mmlu/acc 0.2343, gate PASS. New configuration: 22.4 to 24.0 s wall, mmlu_abstract_algebra/acc 0.27, gate PASS. Also ran ruff check, ruff format --check and pytest --collect-only on the file. A GPU CI run is still needed to confirm the accuracy with CUDA bf16 numerics. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_qwen3_eval_fp8 passes `mmlu_limit=16`, which huggingface_example.sh forwards as `mmlu.py --limit 16`. mmlu.py applies the limit inside evaluate(), which main() calls once per subject, so 16 means up to 16 x 57 = 912 prompts (the comment next to it assumed 16 in total). Each prompt is a separate synchronous `LLM.generate_text([prompt], 2)` call on a 5-shot prompt of up to MAX_SEQ_LEN (4096) tokens. This is the slowest test in CI (mean 723.7 s; min 661.9 s, median 709.9 s, max 969.8 s over 306 passes) and it hit its 900 s timeout in 11 of 333 lane runs. Its runtime does not depend on the host, so it is the sequential calls and the three TRT-LLM start-ups, not the machine. Use `mmlu_limit=1`: 57 prompts, one per subject, so every subject CSV is still read, prompted and scored. Nothing asserts on the MMLU accuracy. That removes 855 synchronous calls: at least 40 s if a call takes 47 ms, and 85 to 255 s at 100 to 300 ms (the per-call latency has not been measured). The lm_eval and simple_eval stages keep their limits, which lm-eval applies per task. Testing: ran the real mmlu.evaluate() with a stub model over three subjects: limit=16 made 16 calls per subject (912 for 57 subjects), limit=1 made one per subject (57). Also ran ruff check, ruff format --check and pytest --collect-only on the file. TRT-LLM and a GPU are needed to measure the saving, so a GPU CI run is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_qwen3_qat_nvfp4 (fsdp2, deepspeed, ddp) and test_qwen3_qad_nvfp4 (fsdp2, deepspeed) each ran quantize.py on the same tiny Qwen3 with the same recipe, and every quantize.py and train.py call built its own tokenized dataset cache (the cache key contains max_length, 8192 in quantize.py and 128 in train.py, so even the two steps of one test could not share it). With test_qwen3_qlora_nvfp4 that is 6 quantize.py runs and 12 cold dataset builds per module. These tests take 104 to 121 s each on CI (qlora 169 s). Run quantize.py once in a module-scoped fixture and let the five tests read the checkpoint, guarded by assert_unmodified_tree like the shared tiny-model fixtures. qat_nvfp4.yaml and qad_nvfp4.yaml give quantize.py identical arguments (distill and teacher_model are not quantize arguments), so one checkpoint serves both. test_qwen3_qlora_nvfp4 keeps its own PTQ because it passes --compress True. All quantize.py and train.py calls now share one dataset cache directory, so the 8192 and 128 variants are each built once; calibration keeps its 8192 tokens. test_dataset_utils_pretokenize keeps its own cache directory so it still checks that the CLI writes one. Fixture time is outside pytest-timeout (timeout_func_only = true), as for the other shared fixtures in tests/examples. Testing: node IDs are unchanged. With run_example_command faked, quantize.py runs 6 -> 2 times and cold dataset builds go 12 -> 2 (plus the pretokenize test's own); a train.py that writes into the shared checkpoint is reported by the fixture teardown. Real quantize.py and train.py (QAT, then QAD from the checkpoint made with the QAT config) ran on CPU with the int4 and fp8 recipes, since NVFP4 dynamic-block quantization needs CUDA: train.py left the checkpoint unmodified and the QAD run reused the dataset cache the QAT run built. The saving, about 3 minutes of the 21.7 minute lane, is an estimate from the CI timings; a GPU CI run is needed to measure it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_llama_eagle3[1-False] is the one compiled EAGLE3 training run on the 1-GPU PR lane: hf_eagle.py compiles _eagle_forward with mode="max-autotune". It takes 175.5 s per job against 40.6 s for [1-True], the same flow without compile (call time 158.5 s vs 34.7 s on PRs, 147.8 s vs 34.4 s nightly), so about 135 s is torch.compile. That includes benchmarking Triton GEMM templates for a draft model that is one tiny decoder layer. Launch it with TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=ATEN, the mechanism test_diffusers.py already uses for its compile tests. The mode stays max-autotune, so cudagraphs, coordinate-descent tuning, Dynamo tracing and Inductor lowering still run; only the Triton GEMM template search is dropped. There is no conv variable here because the draft has no convolutions. [2-False] (2 GPUs, nightly) keeps the full autotune. Testing: with run_example_command faked, only [1-False] gets the variable (on top of the full os.environ) and the other parametrizations pass no env. On torch 2.14 torch._inductor.config.max_autotune_gemm_backends reads the variable and the max-autotune mode options do not override it. CPU cannot reproduce the Triton autotune, so the saving is not measured; a GPU CI run is needed to see how much of the 135 s it removes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_offline_eagle3_training[tinyllama] trains one step (num_train_epochs=0.1) of offline EAGLE3 on the tiny llama with its real base model. It takes 34.9 s per job, almost all launcher start-up. test_eagle_offline_ptq.py's test_offline_eagle_training (35.2 s) already runs that mode on the same model (use_fake_base_for_offline defaults to false) with hidden states dumped by compute_hidden_states_hf.py, and the FakeBaseModel case keeps the synthetic dump path. The only thing the dropped case covered that stage 2 did not is the offline collator truncating dumps: its 128-token dumps were cut to training_seq_len=64, while stage 2's dumps are 28 tokens against 64, so they were only padded. Set stage 2 to training_seq_len=16 so the collator truncates there. The collator is also unit tested on CPU (test_collator_truncates, test_collator_pads). Stage 3 (hf_ptq.py --specdec_offline_dataset) takes its collator length from --calib_seq and does not read the tokenizer saved with the checkpoint, which now records model_max_length 16. Testing: ran stage 1 (compute_hidden_states_hf.py, two 28-token dumps) and stage 2 (main.py, offline EAGLE3 on the tiny llama) on CPU with training_seq_len=64 and 16 (use_cpu, bf16 and tf32 off are the only changes from the test): both train and write the checkpoint, and the collator pads 28 -> 64 and truncates 28 -> 16. With run_example_command faked, the launch commands of the remaining parametrizations are unchanged and the only removed test ID is test_offline_eagle3_training[tinyllama]; nothing else reads its eagle-tinyllama-offline output. The ~30 s saving per job is not measured. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_get_dataset_dataloader_nemotron_v3_chat_template streams the seven gated Nemotron post-training v3 datasets from the Hub and applies a chat template: a 24 s call in the PR run (run 37703177445) that needs an HF token and the network, which also makes it a source of Hub flakes. It checks the live gated data, not modelopt logic: the combo expansion and sample distribution of get_dataset_dataloader are covered by the CPU unit tests in tests/unit/torch/utils/test_dataset_utils.py, which record the per-dataset calls instead of fetching. Mark it `release`. The nightly and on-demand gpu runs pass --run-release, so it still runs there every day, and it is skipped on pull requests. Testing: ruff check and format with the pinned ruff 0.15.20. Without the flag the test is skipped with "Only run when --run-release is given". Not run: the test itself or the nightly. Expected saving is about 24 s per pull-request gpu job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_calibration_data builds its DataLoaders over lists of tensors that are
already on the GPU, so the four workers load nothing. They only add cost: the
script forced the "spawn" start method (CUDA cannot be used in forked workers),
so every loader pass spawned four processes that re-imported torch, timm and
modelopt and received the dataset over CUDA IPC. That is about 14 s per pass
and about 300 s per torch_onnx CI job.
Use num_workers=0 in both loaders and drop the module-level
mp.set_start_method("spawn", force=True). The latter 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; they now run with the default start method,
exactly as in the onnx_ptq example that calls the same function.
Testing: py_compile and ruff 0.15.20 (check and format) on the script, and read
the file for other users of mp and of the start method (none). Not run:
tests/examples/torch_onnx needs a GPU, TensorRT and model downloads, so GPU CI
validation is still needed.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Two cuts to the torch_onnx example lane. swinv2_tiny runs all five formats on top of swin_tiny, which already covers every format through the same exporter and TensorRT path; swinv2 only adds its cosine attention and post-norm blocks. On PRs (GITHUB_REF under refs/heads/pull-request/) skip swinv2 int8, mxfp8 and auto and keep fp8 and nvfp4, which still exercise those blocks. Scheduled and local runs keep every combination, and test ids are unchanged. swinv2_tiny[auto] alone averages 121 s over 304 CI runs (the 22nd slowest example test); the three together cost about 110 s of every PR run. build_trt_engine let trtexec run its default benchmark after the build (200 ms warm-up, then at least 3 s and 10 iterations), and the output is discarded on success. Run a single iteration instead (--warmUp=0 --duration=0 --iterations=1): the engine still executes once, but about 3 s of timing per engine goes away, roughly 40 s per lane run. --trt_build users see no difference besides it returning sooner. Testing: ruff 0.15.20 (check and format), py_compile and pytest --collect-only on CPU (24 tests, ids unchanged). With the example command stubbed, GITHUB_REF=refs/heads/pull-request/123 skips exactly swinv2 int8, mxfp8 and auto, and any other ref skips none. The two in-process mapping tests that import the script pass on CPU, and build_trt_engine passes the three flags to a fake trtexec. Not run: the lane itself (needs a GPU, TensorRT and model downloads) and the flags against a real trtexec, so GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every example step is a fresh `python script.py` that spends 13-22 s of CI time importing torch, transformers and modelopt before it does any work, and an example test runs several steps: 41 launches in the diffusers lane, 21 in torch_onnx, 11 in llm_qat. Only megatron_bridge avoids them today, with a runner that executes in the pytest process and needs the get_args()/main() convention. ExampleRunner runs `python <script>.py ...` steps in one long-lived spawned worker. A step runs as python would run it (runpy as __main__, sys.argv, the example dir as cwd and first on sys.path, output captured at the fd level), and afterwards everything it changed is put back: argv, cwd, sys.path, os.environ, signal handlers, logging handlers and levels, warning filters, RNGs, torch backend flags, the multiprocessing start method, and the modules imported from examples/, so the `utils` or `config` of different examples do not collide. Third-party and modelopt imports stay cached, which is the saving: on a CPU box a step that imports timm, transformers and modelopt.torch.quantization takes 9.4 s as a subprocess, and 0.35 s in the worker after the first one (10.2 s). Per-lane hooks wrap steps for registries the generic restore cannot know about. The worker is a separate process, not the pytest process. A hung step is killed with its children and fails with the output it had printed, and the next step gets a fresh worker. The same happens when pytest-timeout interrupts the wait. A step that takes its worker down fails alone, the worker is recycled after any failed step, and pytest holds no CUDA context. close() lets the worker exit so that coverage writes its data. run_example_command hands an installed runner the steps it accepts (an allowlist of scripts) together with the step's environment; every other step runs as before. StepFailed subclasses CalledProcessError, so callers see what a subprocess would raise, and the Hub retry loop reads its captured output. Plain callables such as the megatron_bridge runner keep their contract. MODELOPT_EXAMPLE_RUNNER=subprocess runs every step as a subprocess, to bisect a suspected leak between steps; =inprocess runs them in the calling process so that pdb works. No lane installs the runner yet. Testing: tests/unit/test_example_runner.py (32 CPU tests, about 10 s) and tests/unit/test_example_run_command.py pass, also under pytest-xdist (-n 4 --dist loadfile) and under COVERAGE_PROCESS_START with COVERAGE_CORE=sysmon (about 30 s). Dropping each of 19 behaviours of the runner (every restore, the process-group kill, the daemon flag, the pipe close, the graceful close) fails at least one test. ruff 0.15.20 and typos are clean. Not run: any GPU lane, so memory released after a step and CUDA RNG state are unvalidated and GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The lane starts `python torch_quant_to_onnx.py` 19 times per PR run, and every start spends 15-20 s importing torch, timm, datasets and modelopt before the step quantizes a one-sample model. Install an ExampleRunner for the lane in tests/examples/torch_onnx/conftest.py, per test as the megatron_bridge conftest does, so that one session-wide worker pays the imports once. The plan estimates -140 to -400 s per PR run. test_hf_embedding_quant_to_onnx.py is not on the runner's list and keeps running as a subprocess, and every step does when MODELOPT_EXAMPLE_RUNNER=subprocess. What the script changes in the process, and what puts it back: sys.path gets ../onnx_ptq, which is restored with the rest of sys.path; the sibling modules evaluation and download_example_onnx are imported from examples/ and dropped after each step; the multiprocessing start method is no longer forced (previous commits) and would be restored anyway; the seeds and cudnn flags of evaluation.py are restored with the other RNG and backend state. The ResNet recipe also adds `_add_resnet_residual_quantizers` to modelopt's CUSTOM_POST_CONVERSION_PLUGINS, a registry that would then apply to every later step. The new keep_post_conversion_plugins hook restores it after each step. Testing: ruff 0.15.20 and typos are clean, and tests/examples/torch_onnx still collects 24 tests on CPU. With the conftest active and CUDA hidden, test_torch_onnx_recipe_flag and test_torch_onnx_auto_quantize_recipe pass through the worker (11.7 s for the first step, 1.3 s for the second) and so do the two mapping tests. The real script through the worker gives the same ONNX op histogram as a subprocess for vit_tiny int8 and fp8: 12.9 s for the first step and 0.6 s for each later one, against 11.7 s per subprocess. A ResNet fp8 recipe step followed by a probe shows the plugin registry restored with the hook and leaked without it, and a unit test of the hook fails when its restore is removed. Not run: anything on a GPU, TensorRT or the --trt_build steps (trtexec runs as a child of the worker), so GPU memory across steps in one process, the worker's coverage data and the lane's timings are unvalidated and GPU CI validation is still needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The Windows unit job failed at collection: the parametrize decorator of
test_worker_death_fails_the_step_not_the_session evaluates -signal.SIGKILL when the module is
imported, and signal has no SIGKILL on Windows. The module's skip_on_windows fixture only runs
at test setup, so the AttributeError aborted the whole session ("Interrupted: 1 error during
collection") and no Windows test ran at all (PR #2730, Unit tests / windows).
Take the parameter from a module-level getattr(signal, "SIGKILL", 9). The script text that sends
the signal still names signal.SIGKILL; it only runs where the test runs, which is POSIX.
Testing: collecting the committed version with signal.SIGKILL, SIGHUP, SIGUSR1 and os.setsid,
killpg and getpgid removed reproduces the CI error. The fixed module, test_distributed_worker_pool.py
and test_example_run_command.py collect (41 tests) under the same conditions. Not run on Windows.
The collection error hid every other Windows result, so the job needs a re-run to show what else,
if anything, fails there.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_search_constraints replaced ConstraintsFunc._get_true_latency with a fake by plain assignment and never put the original back, so the fake stayed in place for the rest of the process. test_evaluate_constraints.py::test_evaluate_latency_constraints and ::test_percent_limits use the real method through a fake runtime client (latency 100 and 200) and fail when they run after it: the leaked fake returns 280.5 for the tiny resnet, so the 20% latency limit comes out as 56.1 instead of 40, and a 200 limit is not met. Serial runs never saw it, because test_evaluate_constraints.py sorts before test_full_algorithms_vision.py. With pytest-xdist (-n 4 --dist loadfile) a worker can get the files in the other order. In PR #2730 the same worker ran test_search_constraints and then the two victims in the Python 3.10 and 3.14 legs, which failed; in the legs where they landed on different workers, or in the other order, they passed. Use monkeypatch, so the patch is undone after the test. Testing: running test_search_constraints[False-True] and then test_evaluate_constraints.py in one process failed both tests with the CI values before the change (56.10 against 40.0, and assert False is True) and passes after it, and ConstraintsFunc._get_true_latency is the original function again afterwards. The whole unit suite in reversed and in shuffled file order is run next to look for other leaks of this kind. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
tests/unit/torch/quantization/test_dist.py::test_data_parallel hit the 60 s default cap in the Python 3.12 / torch 2.13 leg of PR #2730. The test spawns two processes and takes 8-15 s in serial runs. With pytest-xdist it takes 25-35 s in every other leg, and 60 s or more on the slowest runner: four workers and the child processes of the multi-process tests (test_dist, the three test_*data_parallel* autoquant tests, the worker-pool tests) share the four vCPUs of a CI runner. A hung test is still reported, after 2 min instead of 1. Windows already had 120 s. Testing: the durations above come from the job logs of the PR run and of the serial run before it. Not re-run on a runner: the cap only adds headroom. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/ci-speed
branch
from
October 17, 2026 00:38
524d784 to
954ac75
Compare
Collaborator
Author
|
/ok to test 954ac75 |
This was referenced Oct 17, 2026
This branch has not been deployed
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.