Repository navigation
[Draft, not for merge] GPU CI speed-up series (runtime measurement) - #2730
Draft
shengliangxu wants to merge 35 commits into
Draft
shengliangxu wants to merge 35 commits into
shengliangxu wants to merge 35 commits into
Conversation
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). - pyproject.toml: add pytest-xdist to the dev-test extra. uv.lock is refreshed by the weekly bump_uv_lock workflow. - 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. Testing: the full tests/unit suite, on 4 CPUs with coverage as in CI. 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. Not yet run on a real CI runner. 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>
One stuck job takes the rest of a pool-based lane down with it. On 15 of 15 nights the first pytest-timeout in the 2-GPU nightly `gpu` job is test_fsdp2_streaming_export_matches_reference[nvfp4]; the shared session pool never recovers, so 7-12 more `dist_workers` tests then wait out their 120 s call cap, and the job is killed at its 60-minute limit with 667-693 of about 1,700 tests done (39-40%). The latest nightly (run 37865638131) shows the same. Because the job never succeeds, main is never seeded for the extension caches either. Why the pool stays broken: pytest-timeout raises in the pytest process while the ranks are still inside the job, possibly stuck in a collective, and the results of the abandoned job can still arrive in the result queue and be read as the answer to the next test. A worker that died (an illegal memory access can abort the process) was only noticed after the 600 s result timeout. DistributedWorkerPool now: - marks itself broken on any BaseException while waiting for results (pytest-timeout's Failed included), on a timeout, and on a dead worker, and respawns on the next run(). A job that merely raises is reported as before and the workers are reused; - notices a dead worker within a second instead of after 600 s, also while waiting for the workers to start; - gives up with a clear error after two failed restarts in a row, so a broken environment fails the remaining tests at once instead of waiting out the ready timeout for each of them; - exits cleanly (so workers can flush their coverage data) only when healthy, and terminates, then kills, stragglers otherwise; - takes ready_timeout (300 s) and run_timeout (600 s) arguments. Two helpers for pools that outlive test modules (used by the Megatron pool in a following 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), and 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. Testing: new tests/unit/test_distributed_worker_pool.py runs a two-process gloo pool on CPU: a job error keeps the workers (same pids); a timeout respawns them (old pids gone, new ones run); a BaseException raised by a signal handler, as pytest-timeout does, respawns them; a worker that exits fails the run in under 20 s and the pool respawns; repeated start failures end in a fast failure with no new processes; the group cleanup and the flag reset work. 7 tests in 30 s. ruff 0.15.20 clean. Not covered here, and still needing GPU CI: NCCL, CUDA poisoning, real Megatron jobs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_fsdp2_streaming_export_matches_reference[nvfp4] needs two GPUs, so it only runs in the nightly gpu job, and there it hangs until its 120 s call cap on 15 of 15 nights (first timeout at 00:54-01:31 UTC; run 37865638131 on 10-09 included). The root cause has not been found. Before the worker pool respawned after a timeout, the stuck job also took every later `dist_workers` test down with it (7-12 more timeouts) and the job was cancelled at 60 minutes with about 40% of its tests run. The pool now recovers, so the hang would cost 120 s and a respawn instead, but it still ends the test without a result every night. The test has never passed in the nightly in that period, so skipping the nvfp4 parameter loses no coverage that exists today. The fp8 parameter keeps running. Remove the skip when the hang is fixed. Testing: ruff check and format with the pinned 0.15.20, py_compile. The module imports CUDA at collection, so it cannot be collected on the CPU machine used here, and it was not run on two GPUs. 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 against a 120 s ready
timeout.
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. 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;
- _get_pool, when a module first uses a pool: 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;
- DistributedWorkerPool (previous change): a job that times out, is interrupted
or loses a worker makes the pool respawn, which 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 and py_compile, and the pool, group-cleanup and flag-reset
behaviour is covered by the CPU tests of the previous change. This needs Megatron
and GPUs, so it has NOT been run: validate with three consecutive 1-GPU PR runs
and one 2-GPU nightly before relying on it. If a module turns out to depend on a
fresh process, give that module its own pool rather than reverting the rest.
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 one module-scoped vLLM engine per tiny architecture. The dense Llama engine adds almost nothing over the others: the Qwen3-MoE engine already covers QKV and Row parallel linears, base Attention and FusedMoE, and the DeepSeek-V3 engine has a shared expert, so its MergedColumnParallelLinear (gate_up_proj) and RowParallelLinear (down_proj, plus o_proj) are quantized there too. Remove the Llama fixture and test and move its two unique counts, 2 Merged and 4 Row, into test_tiny_deepseek_mla_quantize. DeepSeek's fused_qkv_a_proj has its own forward, so it is not converted and does not contribute to those counts. The only assertion that is dropped is `moe_count == 0`, the negative check that nothing in a dense model is converted as a FusedMoE. Also pass num_gpu_blocks_override=32 and max_num_batched_tokens=64 to the three MoE engines (Qwen3-MoE, DeepSeek-V3, DeepSeek-V3 fp8). 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. The sparse, GLM and DeepSeek-V4 engines are unchanged. Estimated saving: about 28 s per vLLM version for the removed engine and 1-3 s per remaining engine for the smaller pool. Part of the 28 s is first-engine warm-up (kernel compilation and cache population) that the next engine now pays, so the realized saving may be smaller. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. 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. Not run: pytest, because this machine has no GPU or vLLM. GPU CI has to confirm that the engines boot with the smaller pool and that the DeepSeek counts reach 2 Merged and 4 Row. 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>
|
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 71.67% 80.10% +8.42%
==========================================
Files 644 632 -12
Lines 71470 68498 -2972
==========================================
+ Hits 51225 54869 +3644
+ Misses 20245 13629 -6616
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:
|
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>
Collaborator
Author
|
/ok to test 524d784 |
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 (32 commits, 192 added lines of non-test source, no
modelopt/library code touched) so that real runner runs can measure it before it is split into reviewable[x/N]PRs. Those PRs will be linked here.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 series changes:
coverage.xmland usessysmonon the unit, GPU and regression lanes; JIT-extension and Triton caches are kept per lane and saved even when the run fails; a wheel cache for the CUDA packages thegpulane compiles (saved frommainonly); nightly and on-demand runs pass--run-release.DistributedWorkerPoolrespawns after a hang, a timeout or a dead worker (7 new CPU tests); gpu_megatron shares one pool across its modules; longer timeouts for the Mamba test and the pool start; the FSDP2 streaming-export nvfp4 test that hangs the 2-GPU nightly is skipped.release(nightly only).ExampleRunnerworker so example steps share one Python launch (MODELOPT_EXAMPLE_RUNNER=subprocessswitches it off).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
Run locally, CPU only: ruff 0.15.20 and the pre-commit hooks on every commit; actionlint on all workflows; the 7 pool tests, the 33 runner tests, all of
tests/unit/onnx(756 passed, 12 xfailed) and the moved ONNX tests.Nothing in this series has run on a GPU yet; this PR is the first run. What to look at in the runs:
linuxjob time (xdist), which gates every GPU lane;gpulane: the first run is cold, and the JIT and Triton caches fill perpull-request/Nbranch, so a re-run shows warm numbers. The wheel cache can only be seeded frommain, so the three CUDA packages still compile in this PR;--run-releasetests) only run on the nightly or a manual dispatch.Before your PR is "Ready for review"
Draft, not for merge. Commits are signed (
git commit -s -S).CONTRIBUTING.md:pytest-xdistis added to thedev-testextra only;uv.lockis left to the weekly bumpAdditional Information
Planned split into reviewable PRs: CI plumbing (with
--run-release), test infrastructure, per-test trims, and the torch_onnx lane.