From 6bf4ca4157c8a0c1e7b80450eb38c4e50a535ad9 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 18:47:39 +0000
Subject: [PATCH 01/34] Run the unit tests in parallel with pytest-xdist

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>
---
 .github/workflows/unit_tests.yml                        | 7 +++++++
 noxfile.py                                              | 3 ++-
 tests/_test_utils/torch/quantization/onnx_export.py     | 2 +-
 tests/unit/torch/speculative/plugins/test_hf_dflash2.py | 1 +
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index 6fbca185d9f..c3bae2e41be 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -69,6 +69,10 @@ jobs:
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
           # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS).
           PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }}
+          # One single-threaded xdist worker per vCPU (`-n auto` counts physical cores); torch's
+          # default threads per worker oversubscribe the runner and run slower than serial.
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" -- $PYTEST_EXTRA
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v7
@@ -146,6 +150,9 @@ jobs:
         with:
           python-version: ${{ matrix.python_version }}
       - name: Run unit tests
+        env:
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "${{ matrix.nox_session }}"
   partial-install:
     if: needs.check-file-changes.outputs.any_changed == 'true'
diff --git a/noxfile.py b/noxfile.py
index 27d65b26810..15f6d2dc4ae 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -77,11 +77,12 @@ def unit(session, torch_ver, tf_ver):
     if run_manual:
         # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global
         # state that breaks later tests. No coverage, so it cannot overwrite the main run's data.
+        # Serial too (CI's PYTEST_ADDOPTS asks for xdist): a worker would run them in some order.
         session.run(
             "python",
             # Caller args first: the last ``-m`` wins, and it must be ours.
             *("-m", "pytest", "tests/unit", *pytest_args, "--run-manual", "-m", "manual"),
-            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": ""},
+            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": "", "PYTEST_ADDOPTS": ""},
         )
 
 
diff --git a/tests/_test_utils/torch/quantization/onnx_export.py b/tests/_test_utils/torch/quantization/onnx_export.py
index e65e6f23500..b3041f780d1 100644
--- a/tests/_test_utils/torch/quantization/onnx_export.py
+++ b/tests/_test_utils/torch/quantization/onnx_export.py
@@ -23,7 +23,7 @@
 
 from .models import SimpleConv, SimpleConvLinear, SimpleLinear
 
-TEST_MODELS = {SimpleLinear, SimpleConv, SimpleConvLinear}
+TEST_MODELS = (SimpleLinear, SimpleConv, SimpleConvLinear)
 
 
 def onnx_export_tester(model, device, num_bits, per_channel_quantization, constant_folding, dtype):
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
index 931a4a3ee46..a5c8410763c 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
@@ -349,6 +349,7 @@ def test_selector_loss_increases_total_loss(self):
             losses[alpha] = float(out.loss.detach())
         assert losses[1.0] > losses[0.0]
 
+    @pytest.mark.timeout(300)  # ~53 s with torch's default threads, longer single-threaded
     def test_overfits_a_single_batch(self):
         """A few steps on one batch drive backbone and selector accuracy up.
 

From 30d9973c8f8d542519c23df3f9ef8af20f830c0b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:11:44 +0000
Subject: [PATCH 02/34] Write only coverage.xml in CI

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>
---
 pyproject.toml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8843bcfc23e..c98129247cb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -362,8 +362,9 @@ skips = [
 [tool.pytest.ini_options]
 # Default additional options
 # Show a short test summary info for all except passed tests with -ra flag
-# print execution time for 50 slowest tests and generate coverage reports
-addopts = "-v -ra --instafail --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
+# print execution time for 50 slowest tests and write the coverage.xml that CI uploads
+# (add --cov-report=html or --cov-report=term-missing locally; each adds 5-20 s to the end of a run)
+addopts = "-v -ra --instafail --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
 pythonpath = ["tests/"]
 # Apply per-test timeouts (see tests/conftest.py) to the test call only, not fixture setup/teardown
 timeout_func_only = true

From 832b5b5afbfabc776f7fe6b972501aa21570b769 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:12:13 +0000
Subject: [PATCH 03/34] Use the sysmon coverage core on the unit, GPU and
 regression lanes

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>
---
 .github/workflows/gpu_tests.yml        | 3 +++
 .github/workflows/regression_tests.yml | 3 +++
 .github/workflows/unit_tests.yml       | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index a0389da7b62..3bf46d59424 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -88,6 +88,9 @@ jobs:
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the
+          # overhead on import-heavy worker processes. Needs Python 3.12+.
+          COVERAGE_CORE: sysmon
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
index 001c4e6af4a..37f3524afce 100644
--- a/.github/workflows/regression_tests.yml
+++ b/.github/workflows/regression_tests.yml
@@ -60,6 +60,9 @@ jobs:
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the
+          # overhead on import-heavy worker processes. Needs Python 3.12+.
+          COVERAGE_CORE: sysmon
         run: python -m pip install nox && nox -s regression
       - name: Upload regression coverage to Codecov
         uses: codecov/codecov-action@v7
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index c3bae2e41be..3442e6147c1 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -69,6 +69,8 @@ jobs:
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
           # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS).
           PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }}
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the overhead
+          COVERAGE_CORE: sysmon
           # One single-threaded xdist worker per vCPU (`-n auto` counts physical cores); torch's
           # default threads per worker oversubscribe the runner and run slower than serial.
           PYTEST_ADDOPTS: -n 4 --dist loadfile

From 7f8854a8783821e95dabac5efb1b8145ab56bec4 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:12:54 +0000
Subject: [PATCH 04/34] Cache JIT extensions per lane, save them on failure and
 cache Triton 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>
---
 .github/actions/cache-extensions/action.yml   | 44 +++++++++++++++----
 .../actions/cache-extensions/save/action.yml  | 31 +++++++++++++
 .github/workflows/_example_tests_runner.yml   | 16 ++++++-
 .github/workflows/gpu_tests.yml               | 16 ++++++-
 4 files changed, 95 insertions(+), 12 deletions(-)
 create mode 100644 .github/actions/cache-extensions/save/action.yml

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index 7ddf65de8ce..3043621f180 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -1,24 +1,52 @@
 name: Cache extensions
-description: Cache and reuse JIT-compiled extensions (e.g. CUDA extensions) across runs.
+description: Restore JIT-compiled CUDA extensions and Triton kernels cached by earlier runs. Pair with cache-extensions/save after the tests.
 
 inputs:
   cache-key:
-    description: Environment discriminator for the cache key (e.g. GPU arch + container image).
+    description: Environment discriminator for the cache keys (e.g. GPU arch + lane + container image).
     required: true
 
+outputs:
+  ext-key:
+    description: Key under which cache-extensions/save stores the torch extensions cache.
+    value: ${{ steps.ext.outputs.cache-primary-key }}
+  ext-hit:
+    description: "'true' if the torch extensions cache was restored with an exact key match."
+    value: ${{ steps.ext.outputs.cache-hit }}
+  triton-key:
+    description: Key under which cache-extensions/save stores the Triton cache.
+    value: ${{ steps.triton.outputs.cache-primary-key }}
+  triton-hit:
+    description: "'true' if the Triton cache was restored with an exact key match."
+    value: ${{ steps.triton.outputs.cache-hit }}
+
 runs:
   using: composite
   steps:
     - shell: bash
-      run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
-    - id: cache
-      uses: actions/cache@v6
+      run: |
+        echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
+        echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV"
+    # CUDA extensions: exact key only, because ninja decides what to rebuild from file timestamps.
+    # Only CUDA/C++ sources and the loaders are hashed, so Triton-only edits keep the entry valid.
+    - id: ext
+      uses: actions/cache/restore@v6
       with:
         path: /root/.cache/torch_extensions
-        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**', 'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py')
-          }}
+        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cu', 'modelopt/torch/kernels/quantization/**/*.cuh',
+          'modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.h', 'modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py',
+          'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py') }}
     # On a cache hit, backdate sources below the cached objects so ninja reuses them (touching
     # the objects instead would desync ninja's .ninja_deps and force a rebuild).
-    - if: steps.cache.outputs.cache-hit == 'true'
+    - if: steps.ext.outputs.cache-hit == 'true'
       shell: bash
       run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} +
+    # Triton: entries are content-addressed, so restoring a cache from an older commit is safe
+    # and only costs the kernels that changed.
+    - id: triton
+      uses: actions/cache/restore@v6
+      with:
+        path: /root/.cache/triton
+        key: triton-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }}
+        restore-keys: |
+          triton-${{ inputs.cache-key }}-
diff --git a/.github/actions/cache-extensions/save/action.yml b/.github/actions/cache-extensions/save/action.yml
new file mode 100644
index 00000000000..ed54b18853c
--- /dev/null
+++ b/.github/actions/cache-extensions/save/action.yml
@@ -0,0 +1,31 @@
+name: Save extension caches
+description: Save the JIT caches that cache-extensions did not restore with an exact key match. Run it under always() so a failed or cancelled run still
+  seeds the next one.
+
+inputs:
+  ext-key:
+    description: The ext-key output of cache-extensions.
+    required: true
+  ext-hit:
+    description: The ext-hit output of cache-extensions.
+    required: true
+  triton-key:
+    description: The triton-key output of cache-extensions.
+    required: true
+  triton-hit:
+    description: The triton-hit output of cache-extensions.
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - if: inputs.ext-hit != 'true'
+      uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/torch_extensions
+        key: ${{ inputs.ext-key }}
+    - if: inputs.triton-hit != 'true'
+      uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/triton
+        key: ${{ inputs.triton-key }}
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index e3378865b65..22664c224d8 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -58,9 +58,12 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
-      - uses: ./.github/actions/cache-extensions
+      - id: jit-cache
+        uses: ./.github/actions/cache-extensions
         with:
-          cache-key: rtxpro6000-${{ inputs.docker_image }}
+          # One entry per example: the GPU-tests lanes and the examples that share an image would
+          # otherwise share one immutable, partially filled entry.
+          cache-key: rtxpro6000-example-${{ inputs.example }}-${{ inputs.docker_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
@@ -113,6 +116,15 @@ jobs:
             pytest_args+=(-n 4)
           fi
           python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}"
+      - name: Save JIT caches
+        # Also after a failed or cancelled run, so the next run starts warm
+        if: always()
+        uses: ./.github/actions/cache-extensions/save
+        with:
+          ext-key: ${{ steps.jit-cache.outputs.ext-key }}
+          ext-hit: ${{ steps.jit-cache.outputs.ext-hit }}
+          triton-key: ${{ steps.jit-cache.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-cache.outputs.triton-hit }}
       - name: Flag allowed failure
         if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }}
         run: |
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 3bf46d59424..e5dd45b583d 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -78,9 +78,12 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
-      - uses: ./.github/actions/cache-extensions
+      - id: jit-cache
+        uses: ./.github/actions/cache-extensions
         with:
-          cache-key: rtxpro6000-${{ matrix.container_image }}
+          # One entry per lane: lanes sharing an image (gpu_megatron and the megatron_bridge
+          # example) would otherwise share one immutable, partially filled entry.
+          cache-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
@@ -94,6 +97,15 @@ jobs:
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
+      - name: Save JIT caches
+        # Also after a failed or cancelled run, so the next run starts warm
+        if: always()
+        uses: ./.github/actions/cache-extensions/save
+        with:
+          ext-key: ${{ steps.jit-cache.outputs.ext-key }}
+          ext-hit: ${{ steps.jit-cache.outputs.ext-hit }}
+          triton-key: ${{ steps.jit-cache.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-cache.outputs.triton-hit }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v7
         with:

From a355d34f178dac8c9be04efa4945143d0dd74a47 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:13:12 +0000
Subject: [PATCH 05/34] Give the Mamba search-space test and the Megatron
 worker pool more time

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>
---
 tests/_test_utils/torch/distributed/utils.py                   | 3 ++-
 .../torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py  | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py
index ea954efbcf5..dd8123d03f9 100644
--- a/tests/_test_utils/torch/distributed/utils.py
+++ b/tests/_test_utils/torch/distributed/utils.py
@@ -161,7 +161,8 @@ def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardo
             self._processes.append(p)
 
         for _ in range(world_size):
-            msg = self._result_queue.get(timeout=120)
+            # Cold imports of the Megatron/TE/modelopt stack take 75-100 s under coverage in CI
+            msg = self._result_queue.get(timeout=300)
             assert msg == "ready", f"Worker failed to initialize: {msg}"
 
     @staticmethod
diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
index 8c98f7127fa..32df35a719c 100644
--- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
+++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 
+import pytest
 import torch
 from _test_utils.import_helper import skip_if_no_mamba
 
@@ -133,6 +134,7 @@ def _test_mamba_search_space(rank, size):
     assert not any(named_dynamic_modules(model))
 
 
+@pytest.mark.timeout(240)  # The first Mamba (Triton) kernel compile takes 75-100+ s on a cold cache
 def test_mamba_search_space(dist_workers):
     dist_workers.run(_test_mamba_search_space)
 

From 7faaa6860830ab9d6bda2a2083e133d8cb48c5b6 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 18:51:31 +0000
Subject: [PATCH 06/34] Notice a dead pool worker within a second and add
 worker-state helpers

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>
---
 tests/_test_utils/torch/distributed/utils.py |  90 ++++++++-
 tests/unit/test_distributed_worker_pool.py   | 185 +++++++++++++++++++
 2 files changed, 269 insertions(+), 6 deletions(-)
 create mode 100644 tests/unit/test_distributed_worker_pool.py

diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py
index dd8123d03f9..515b3207e54 100644
--- a/tests/_test_utils/torch/distributed/utils.py
+++ b/tests/_test_utils/torch/distributed/utils.py
@@ -15,6 +15,7 @@
 
 import contextlib
 import faulthandler
+import gc
 import os
 import queue
 import signal
@@ -96,6 +97,53 @@ def get_live_pool(pools, key, make_pool):
     return pool
 
 
+def destroy_extra_process_groups():
+    """Destroy every process group except the default one.
+
+    Libraries that create groups per model (e.g. Megatron's ``initialize_model_parallel``, ~20 NCCL
+    groups) drop their references on teardown but not the groups themselves. In a worker that
+    lives through hundreds of jobs they pile up (communicators, one watchdog thread each).
+    """
+    from torch.distributed.distributed_c10d import _get_default_group, _world
+
+    default = _get_default_group()
+    for group in [g for g in _world.pg_map if g is not default]:
+        try:
+            dist.destroy_process_group(group)
+        except Exception as e:  # noqa: PERF203
+            print(f"Error destroying process group: {e}")
+
+
+_BASELINE_TORCH_FLAGS = {}
+
+
+def _capture_baseline_torch_flags():
+    """Remember the torch settings a freshly spawned worker starts with."""
+    _BASELINE_TORCH_FLAGS.update(
+        cudnn_deterministic=torch.backends.cudnn.deterministic,
+        cudnn_benchmark=torch.backends.cudnn.benchmark,
+        cudnn_tf32=torch.backends.cudnn.allow_tf32,
+        matmul_tf32=torch.backends.cuda.matmul.allow_tf32,
+        deterministic_algorithms=torch.are_deterministic_algorithms_enabled(),
+    )
+
+
+def reset_worker_state(rank, world_size):
+    """Return a persistent worker to the state it had right after spawning.
+
+    ``set_seed`` and similar helpers change process-global torch settings; run this job at the
+    boundary between test modules that share workers.
+    """
+    torch.backends.cudnn.deterministic = _BASELINE_TORCH_FLAGS["cudnn_deterministic"]
+    torch.backends.cudnn.benchmark = _BASELINE_TORCH_FLAGS["cudnn_benchmark"]
+    torch.backends.cudnn.allow_tf32 = _BASELINE_TORCH_FLAGS["cudnn_tf32"]
+    torch.backends.cuda.matmul.allow_tf32 = _BASELINE_TORCH_FLAGS["matmul_tf32"]
+    torch.use_deterministic_algorithms(_BASELINE_TORCH_FLAGS["deterministic_algorithms"])
+    torch.manual_seed(1234)
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
 def default_worker_teardown(rank, world_size):
     """Minimal cleanup between tests in persistent workers."""
     try:
@@ -131,7 +179,9 @@ def _test_fn(rank, size): ...
         pool.shutdown()
     """
 
-    def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardown):
+    def __init__(
+        self, world_size, backend="nccl", teardown_fn=default_worker_teardown, ready_timeout=300
+    ):
         assert world_size > 0, "World size must be greater than 0"
         self.world_size = world_size
         # Set once a job was interrupted: the workers are then killed, see ``run``.
@@ -160,10 +210,37 @@ def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardo
             p.start()
             self._processes.append(p)
 
-        for _ in range(world_size):
-            # Cold imports of the Megatron/TE/modelopt stack take 75-100 s under coverage in CI
-            msg = self._result_queue.get(timeout=300)
-            assert msg == "ready", f"Worker failed to initialize: {msg}"
+        try:
+            for _ in range(world_size):
+                # Cold imports of the Megatron/TE/modelopt stack take 75-100 s under coverage in CI
+                msg = self._get(ready_timeout)
+                assert msg == "ready", f"Worker failed to initialize: {msg}"
+        except BaseException:
+            self._abort()  # do not leave workers behind when the start fails
+            raise
+
+    def _get(self, timeout):
+        """Return the next message from the workers; fail within a second if a worker died.
+
+        A worker that exits (an illegal memory access aborts the process, an OOM kill) would
+        otherwise only be noticed when ``timeout`` runs out.
+        """
+        deadline = time.monotonic() + timeout
+        while True:
+            try:
+                return self._result_queue.get(
+                    timeout=max(0.01, min(1.0, deadline - time.monotonic()))
+                )
+            except queue.Empty:  # noqa: PERF203
+                dead = [
+                    (rank, p.exitcode) for rank, p in enumerate(self._processes) if not p.is_alive()
+                ]
+                if dead:
+                    raise RuntimeError(
+                        f"Worker process(es) died (rank, exit code): {dead}"
+                    ) from None
+                if time.monotonic() >= deadline:
+                    raise
 
     @staticmethod
     def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teardown_fn):
@@ -179,6 +256,7 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard
         if backend == "nccl" and torch.cuda.is_available():
             torch.cuda.set_device(rank)
         torch.manual_seed(1234)
+        _capture_baseline_torch_flags()
         result_queue.put("ready")
 
         while True:
@@ -227,7 +305,7 @@ def run(self, fn, *args, **kwargs):
                 q.put((fn, args, kwargs))
             for _ in range(self.world_size):
                 timeout = self.job_timeout if deadline is None else deadline - time.monotonic()
-                status, rank, tb = self._result_queue.get(timeout=max(timeout, 0))
+                status, rank, tb = self._get(max(timeout, 0))
                 if status == "error":
                     errors.append(f"--- Rank {rank} ---\n{tb}")
                     deadline = deadline or time.monotonic() + self.failure_grace
diff --git a/tests/unit/test_distributed_worker_pool.py b/tests/unit/test_distributed_worker_pool.py
new file mode 100644
index 00000000000..e41bc8e7420
--- /dev/null
+++ b/tests/unit/test_distributed_worker_pool.py
@@ -0,0 +1,185 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU (gloo) tests for what ``DistributedWorkerPool`` adds to its hang handling (see tests/gpu)."""
+
+import multiprocessing
+import os
+import signal
+import time
+from functools import partial
+from pathlib import Path
+
+import pytest
+import torch
+import torch.distributed as dist
+from _test_utils.torch.distributed.utils import (
+    DistributedWorkerPool,
+    destroy_extra_process_groups,
+    reset_worker_state,
+)
+
+pytestmark = [pytest.mark.usefixtures("skip_on_windows"), pytest.mark.timeout(120)]
+
+
+def _touch(rank, size, directory, name="ran"):
+    (Path(directory) / f"{name}{rank}").write_text(str(os.getpid()))
+
+
+def _raise_on_rank1(rank, size):
+    if rank == 1:
+        raise ValueError("boom")
+
+
+def _hang(rank, size):
+    time.sleep(3600)
+
+
+def _die_on_rank0(rank, size):
+    if rank == 0:
+        os._exit(13)
+    time.sleep(3600)
+
+
+def _extra_groups_are_destroyed(rank, size, directory):
+    from torch.distributed.distributed_c10d import _world
+
+    groups = [dist.new_group([0, 1]) for _ in range(3)]
+    dist.all_reduce(torch.ones(1), group=groups[0])
+    assert len(_world.pg_map) == 1 + 3
+
+    destroy_extra_process_groups()
+    assert len(_world.pg_map) == 1
+
+    # Groups can be created and used again afterwards, and the default group still works
+    dist.all_reduce(torch.ones(1), group=dist.new_group([0, 1]))
+    destroy_extra_process_groups()
+    dist.all_reduce(torch.ones(1))
+    (Path(directory) / f"groups{rank}").write_text("ok")
+
+
+def _flags_are_reset(rank, size, directory):
+    baseline = (
+        torch.backends.cudnn.deterministic,
+        torch.backends.cudnn.benchmark,
+        torch.backends.cudnn.allow_tf32,
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.are_deterministic_algorithms_enabled(),
+    )
+    torch.backends.cudnn.deterministic = not baseline[0]
+    torch.backends.cudnn.benchmark = not baseline[1]
+    torch.backends.cudnn.allow_tf32 = not baseline[2]
+    torch.backends.cuda.matmul.allow_tf32 = not baseline[3]
+    torch.use_deterministic_algorithms(not baseline[4])
+    torch.manual_seed(5)
+
+    reset_worker_state(rank, size)
+
+    assert baseline == (
+        torch.backends.cudnn.deterministic,
+        torch.backends.cudnn.benchmark,
+        torch.backends.cudnn.allow_tf32,
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.are_deterministic_algorithms_enabled(),
+    )
+    assert torch.initial_seed() == 1234
+    (Path(directory) / f"flags{rank}").write_text("ok")
+
+
+class _Interrupt(BaseException):
+    """Stands in for pytest-timeout's ``Failed``, which is also a ``BaseException``."""
+
+
+@pytest.fixture
+def pool():
+    pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None, ready_timeout=90)
+    yield pool
+    pool.shutdown()
+
+
+def _pids(pool):
+    return [p.pid for p in pool._processes]
+
+
+def _is_dead(pid):
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return True
+    return False
+
+
+def test_job_error_keeps_the_workers(pool, tmp_path):
+    pids = _pids(pool)
+    pool.run(partial(_touch, directory=tmp_path))
+    with pytest.raises(RuntimeError, match="boom"):
+        pool.run(_raise_on_rank1)
+    pool.run(partial(_touch, directory=tmp_path, name="again"))
+
+    assert not pool.broken
+    assert _pids(pool) == pids
+    assert {int(p.read_text()) for p in tmp_path.glob("again*")} == set(pids)
+
+
+def test_interrupt_kills_the_workers_and_breaks_the_pool(pool):
+    # What pytest-timeout's signal method does: raise in the main thread while the run waits.
+    def interrupt(signum, frame):
+        raise _Interrupt
+
+    pids = _pids(pool)
+    previous = signal.signal(signal.SIGALRM, interrupt)
+    signal.setitimer(signal.ITIMER_REAL, 2)
+    try:
+        with pytest.raises(_Interrupt):
+            pool.run(_hang)
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, previous)
+
+    assert pool.broken
+    assert all(_is_dead(pid) for pid in pids)
+    with pytest.raises(RuntimeError, match="killed"):
+        pool.run(_hang)
+
+
+def test_dead_worker_fails_fast_and_breaks_the_pool(pool):
+    pids = _pids(pool)
+    started = time.monotonic()
+    with pytest.raises(RuntimeError, match=r"died.*13"):
+        pool.run(_die_on_rank0)
+    assert time.monotonic() - started < 20  # not the 600 s job timeout
+
+    assert pool.broken
+    assert all(_is_dead(pid) for pid in pids)  # the survivor, stuck in the job, is killed too
+
+
+def test_failed_start_raises_and_leaves_no_workers():
+    before = {p.pid for p in multiprocessing.active_children()}
+    started = time.monotonic()
+    with pytest.raises(RuntimeError, match="died"):
+        DistributedWorkerPool(world_size=2, backend="not-a-backend", teardown_fn=None)
+    assert time.monotonic() - started < 60  # not the 300 s ready timeout
+    time.sleep(1)
+    assert {p.pid for p in multiprocessing.active_children()} <= before
+
+
+def test_extra_process_groups_are_destroyed(pool, tmp_path):
+    pool.run(partial(_extra_groups_are_destroyed, directory=tmp_path))
+    assert len(list(tmp_path.glob("groups*"))) == 2
+
+
+def test_reset_worker_state_restores_the_spawn_time_settings(pool, tmp_path):
+    pool.run(partial(_flags_are_reset, directory=tmp_path))
+    assert len(list(tmp_path.glob("flags*"))) == 2

From 2daa767b7222a38142dbce6e759210327ede16cc Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 18:55:59 +0000
Subject: [PATCH 07/34] Share one Megatron worker pool across the gpu_megatron
 modules

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>
---
 tests/gpu_megatron/conftest.py | 42 +++++++++++++++++++++++++---------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py
index 5760b9fe22d..8b4490e314e 100644
--- a/tests/gpu_megatron/conftest.py
+++ b/tests/gpu_megatron/conftest.py
@@ -13,11 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
+import gc
 
 import pytest
 import torch
 from _test_utils.fs_utils import assert_unmodified_tree
-from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool
+from _test_utils.torch.distributed.utils import (
+    DistributedWorkerPool,
+    destroy_extra_process_groups,
+    get_live_pool,
+    reset_worker_state,
+)
 from _test_utils.torch.transformers_models import get_tiny_tokenizer
 from megatron.core.parallel_state import destroy_model_parallel
 
@@ -68,6 +74,8 @@ def _prestart_preprocess_forkserver():
 
 def megatron_worker_teardown(rank, world_size):
     """Clean up model-parallel state between tests in persistent workers."""
+    # Surface asynchronous CUDA errors at the test that caused them, not at a later one
+    torch.cuda.synchronize()
     if dist.is_initialized():
         dist.barrier()
     try:
@@ -79,6 +87,9 @@ def megatron_worker_teardown(rank, world_size):
             apex_destroy()
         except Exception as e:
             print(f"Error destroying model parallel with Apex: {e}")
+    # destroy_model_parallel() drops Megatron's references but leaves the NCCL groups behind
+    destroy_extra_process_groups()
+    gc.collect()
     torch.cuda.empty_cache()
 
 
@@ -90,16 +101,17 @@ def _make_pool(world_size):
     )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def _pool_cache():
-    """Module-scoped cache of worker pools keyed by world_size.
+    """Session-scoped cache of worker pools keyed by world_size.
 
     Spinning up a pool cold-imports the full torch/megatron/modelopt stack per worker
-    (~tens of seconds), so fixtures that request the same world_size share one pool
-    instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and
-    ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at
-    module end, so workers are never reused across modules (avoids cross-test
-    state contamination).
+    (75-100 s in CI), so every module that requests the same world_size shares one pool,
+    e.g. on a 2-GPU runner ``dist_workers`` and ``dist_workers_size_2`` are both size 2.
+    Isolation between tests comes from ``megatron_worker_teardown`` (model-parallel state and
+    process groups after every job), from ``_reset_shared_pools`` (process-global torch settings
+    at each module boundary) and from ``get_live_pool`` replacing a pool that a hung, timed-out
+    or crashed job killed.
     """
     pools: dict[int, DistributedWorkerPool] = {}
     yield pools
@@ -107,21 +119,29 @@ def _pool_cache():
         pool.shutdown()
 
 
+@pytest.fixture(scope="module", autouse=True)
+def _reset_shared_pools(_pool_cache):
+    """Workers outlive test modules: undo what the previous module did to process-global torch state."""
+    for pool in _pool_cache.values():
+        if not pool.broken:
+            pool.run(reset_worker_state)
+
+
 def _get_pool(cache, world_size):
     # A hung test kills its pool; a fresh one keeps a single hang from timing out every later
-    # test in the module (that cascade cost the nightly 20+ min).
+    # test (that cascade cost the nightly 20+ min).
     return get_live_pool(cache, world_size, lambda: _make_pool(world_size))
 
 
 @pytest.fixture
 def dist_workers(_pool_cache):
-    """Pool (cached per module) with world_size=torch.cuda.device_count()."""
+    """Shared pool with world_size=torch.cuda.device_count()."""
     return _get_pool(_pool_cache, torch.cuda.device_count())
 
 
 @pytest.fixture
 def dist_workers_size_1(_pool_cache):
-    """Pool (cached per module) with world_size=1 for tests that require a single process."""
+    """Shared pool with world_size=1 for tests that require a single process."""
     return _get_pool(_pool_cache, 1)
 
 

From e71b16625ce2926f96d08a143b2367559347126a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:25:53 +0000
Subject: [PATCH 08/34] Cache the CUDA wheels the gpu lane compiles

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>
---
 .github/workflows/gpu_tests.yml | 26 +++++++++++++++++++
 noxfile.py                      | 46 ++++++++++++++++++++++++++++++---
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index e5dd45b583d..df9e68d6725 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -70,6 +70,8 @@ jobs:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         # Build CUDA kernels only for the runner's RTX PRO 6000 (sm_120), not the image's ~6 archs.
         TORCH_CUDA_ARCH_LIST: "12.0"
+        # Where the gpu lane keeps the CUDA wheels it would otherwise compile (see noxfile.py).
+        WHEELHOUSE: /root/.cache/wheelhouse
     steps:
       - name: Install git
         # The vllm container ships without git; needed for a real checkout (correct
@@ -87,6 +89,30 @@ jobs:
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
+      # The gpu lane compiles mamba_ssm, causal-conv1d and fast-hadamard-transform, 6-8 min. The
+      # nightly run on main builds them once per container image and ISO week (so new releases
+      # are picked up) and saves the wheels; every other run restores them. Without a cached
+      # wheelhouse the gpu session compiles them itself.
+      - name: Get the ISO week
+        id: week
+        if: matrix.example == 'gpu'
+        run: echo "iso=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT"
+      - name: Restore CUDA wheels
+        id: wheelhouse
+        if: matrix.example == 'gpu'
+        uses: actions/cache/restore@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: wheelhouse-${{ matrix.container_image }}-${{ steps.week.outputs.iso }}-${{ hashFiles('noxfile.py') }}
+      - name: Build CUDA wheels
+        if: matrix.example == 'gpu' && github.ref == 'refs/heads/main' && steps.wheelhouse.outputs.cache-hit != 'true'
+        run: python3 -m pip install nox && nox -s gpu_wheels
+      - name: Save CUDA wheels
+        if: matrix.example == 'gpu' && github.ref == 'refs/heads/main' && steps.wheelhouse.outputs.cache-hit != 'true'
+        uses: actions/cache/save@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: ${{ steps.wheelhouse.outputs.cache-primary-key }}
       - name: Run gpu tests
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
diff --git a/noxfile.py b/noxfile.py
index 15f6d2dc4ae..7e82803ed95 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -114,6 +114,25 @@ def partial_unit(session, subset):
 # existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one.
 # Use `python -m pip/pytest` to ensure the container's active venv Python is used,
 # not a stale PATH entry (e.g. NeMo container has pip → /usr/local/bin/pip but python → /opt/venv/bin/python).
+
+# CUDA packages the gpu session compiles against the container's torch, which takes ~7 min together
+# (wheel file prefix -> requirement). CI builds them once with `nox -s gpu_wheels` into $WHEELHOUSE
+# (an Actions cache); the gpu session installs the wheels found there instead of compiling again.
+_WHEELHOUSE = os.environ.get("WHEELHOUSE")
+_CUDA_BUILDS = {
+    "fast_hadamard_transform": "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+    # Latest *released* sdists
+    "mamba_ssm": "mamba_ssm",
+    "causal_conv1d": "causal-conv1d",
+}
+
+
+def _cuda_build(name):
+    """The wheel of ``name`` from the wheelhouse if it has one, else the requirement to build it from."""
+    wheels = sorted(glob.glob(f"{_WHEELHOUSE}/{name}-*.whl")) if _WHEELHOUSE else []
+    return wheels[-1] if wheels else _CUDA_BUILDS[name]
+
+
 # Container: nvcr.io/nvidia/pytorch:26.07-py3 or later with CUDA 13
 @nox.session(venv_backend="none")
 def gpu(session):
@@ -125,7 +144,7 @@ def gpu(session):
         "pip",
         "install",
         "--no-build-isolation",
-        "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+        _cuda_build("fast_hadamard_transform"),
     )
     session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]")
     session.run(
@@ -134,13 +153,32 @@ def gpu(session):
         "pip",
         "install",
         "--no-build-isolation",
-        # Install the latest *released* sdists (built against the container torch)
-        "mamba_ssm",
-        "causal-conv1d",
+        _cuda_build("mamba_ssm"),
+        _cuda_build("causal_conv1d"),
     )
     session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
 
 
+# Container: same as gpu
+@nox.session(venv_backend="none")
+def gpu_wheels(session):
+    """Build the CUDA wheels of the gpu session into $WHEELHOUSE."""
+    if not _WHEELHOUSE:
+        session.error("Set WHEELHOUSE to the directory the wheels are built into")
+    for requirement in _CUDA_BUILDS.values():
+        session.run(
+            "python",
+            "-m",
+            "pip",
+            "wheel",
+            "--no-build-isolation",
+            "--no-deps",
+            "-w",
+            _WHEELHOUSE,
+            requirement,
+        )
+
+
 # Container: nvcr.io/nvidia/nemo:26.08 or later
 @nox.session(venv_backend="none")
 def gpu_megatron(session):

From 15c7255968fa5699d0c175d37f4ed238c0294342 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:43:23 +0000
Subject: [PATCH 09/34] Run the tests marked release on the nightly and
 on-demand gpu runs

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>
---
 .github/workflows/gpu_tests.yml         | 3 +++
 docs/source/deployment/3_unified_hf.rst | 4 ++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index df9e68d6725..c9ae39ebc2a 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -120,6 +120,9 @@ jobs:
           # sys.monitoring records the same lines as the default C tracer at a fraction of the
           # overhead on import-heavy worker processes. Needs Python 3.12+.
           COVERAGE_CORE: sysmon
+          # Tests marked `release` (real-weight downloads and the like) run on nightly and
+          # on-demand runs, not on pull requests.
+          PYTEST_ADDOPTS: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') && '--run-release' || '' }}
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst
index d3dc840cc62..c231169cd23 100644
--- a/docs/source/deployment/3_unified_hf.rst
+++ b/docs/source/deployment/3_unified_hf.rst
@@ -144,8 +144,8 @@ prompts, asserting that each returns non-empty output.
 Two limits are worth stating plainly, because they bound what any ✅ below can mean:
 
 * **These are declared cases, not PR-gated coverage.** The suite is marked ``release`` and collects
-  only when pytest is given ``--run-release``, which no workflow in ``.github/workflows/`` currently
-  passes. A green check on a pull request does not mean these cases ran.
+  only when pytest is given ``--run-release``, which no example workflow in ``.github/workflows/``
+  currently passes. A green check on a pull request does not mean these cases ran.
 * **Each case is a load-and-generate smoke check on the text path.** It does not verify accuracy,
   image or audio inputs, diffusion output, or that speculative decoding actually engages.
 

From 985713a5631ec050db073ec8625af937e56abeea Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:43:45 +0000
Subject: [PATCH 10/34] Run the real-weights Megatron MMLU test only with
 --run-release

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>
---
 tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py b/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py
index 8d3af2d7f3c..286eea97a5b 100644
--- a/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py
+++ b/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py
@@ -64,6 +64,9 @@ def _test_megatron_generate_and_mmlu(rank, size, parallelism):
     assert 0.36 < megatron_mmlu(model, tokenizer, fraction=0.1, batch_size=16) < 0.39
 
 
+@pytest.mark.release(
+    reason="Real Qwen3-0.6B weights and a 1.2-1.5 GB download; run with --run-release."
+)
 @pytest.mark.parametrize("parallelism", ["tp", "pp", "cp", "dp"])
 def test_megatron_generate_and_mmlu(dist_workers, parallelism, num_gpus):
     if num_gpus == 1 and parallelism != "tp":

From 533945e44b901c453aab23ee66fd5656c129a496 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 18:58:39 +0000
Subject: [PATCH 11/34] Size the vLLM MoE engines' KV pool down

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>
---
 .../quantization/test_vllm_dynamic_modules.py | 20 +++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py
index b8a0464b841..6e0fd19dc68 100644
--- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py
+++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py
@@ -1083,7 +1083,13 @@ def _missing(module, name, slots):
     }
 
 
-def _boot_llm(model_dir, max_model_len=64, max_num_seqs=1, **extra):
+_MAX_MODEL_LEN = 64
+# One 64-token request needs a handful of KV blocks. The defaults size the KV pool from the memory
+# fraction (about 17 GB, millions of Python block objects) and the token budget to thousands.
+_SMALL_ENGINE = {"num_gpu_blocks_override": 32, "max_num_batched_tokens": _MAX_MODEL_LEN}
+
+
+def _boot_llm(model_dir, max_model_len=_MAX_MODEL_LEN, max_num_seqs=1, **extra):
     """Construct a vLLM engine on a tiny model.
 
     MoE fixtures override with ``moe_backend="triton"`` (pins the Triton
@@ -1154,7 +1160,9 @@ def _tiny_moe_llm(tmp_path_factory, model_type):
         num_experts_per_tok=2,
         **expert_config,
     )
-    llm = _boot_llm(model_dir, moe_backend="triton", enable_expert_parallel=True)
+    # GPT-OSS has sliding-window layers; the KV pool override is not verified on them.
+    small = _SMALL_ENGINE if model_type == "qwen3_moe" else {}
+    llm = _boot_llm(model_dir, moe_backend="triton", enable_expert_parallel=True, **small)
     try:
         yield llm
     finally:
@@ -1184,7 +1192,7 @@ def tiny_deepseek_llm(tmp_path_factory):
     model_dir = create_tiny_deepseek_v3_dir(
         tmp, kv_lora_rank=512, qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128
     )
-    llm = _boot_llm(model_dir, moe_backend="triton", enable_expert_parallel=True)
+    llm = _boot_llm(model_dir, moe_backend="triton", enable_expert_parallel=True, **_SMALL_ENGINE)
     try:
         yield llm
     finally:
@@ -1782,7 +1790,11 @@ def tiny_deepseek_fp8_llm(tmp_path):
         tmp_path, kv_lora_rank=512, qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128
     )
     llm = _boot_llm(
-        model_dir, quantization="fp8", moe_backend="triton", enable_expert_parallel=True
+        model_dir,
+        quantization="fp8",
+        moe_backend="triton",
+        enable_expert_parallel=True,
+        **_SMALL_ENGINE,
     )
     try:
         yield llm

From 576c8f5af3f398eca84a89470d27537014892cbb Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:44:22 +0000
Subject: [PATCH 12/34] Cut the DFlash regression to 500 training steps and a
 20-step resume

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>
---
 .../regression/torch/speculative/test_dflash.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/tests/regression/torch/speculative/test_dflash.py b/tests/regression/torch/speculative/test_dflash.py
index 18e95148fcd..82a80a17153 100644
--- a/tests/regression/torch/speculative/test_dflash.py
+++ b/tests/regression/torch/speculative/test_dflash.py
@@ -61,6 +61,10 @@
     "dflash.dflash_architecture_config.num_hidden_layers=2",
 ]
 
+# Steps, not epochs: accelerate runs one process per GPU, so an epoch is 500 steps on 1 GPU, 250 on 2.
+_TRAIN_STEPS = 500
+_RESUME_STEPS = 20
+
 
 @pytest.fixture(scope="session")
 def qwen3_model_name():
@@ -80,8 +84,8 @@ def test_dflash_training(qwen3_model_name, dflash_output_dir):
     overrides = [
         f"model.model_name_or_path={qwen3_model_name}",
         f"training.output_dir={output_dir}",
-        "training.num_train_epochs=3",
-        "training.save_steps=500",
+        f"training.max_steps={_TRAIN_STEPS}",
+        f"training.save_steps={_TRAIN_STEPS}",
         *_DFLASH_OVERRIDES,
     ]
 
@@ -108,7 +112,7 @@ def test_dflash_training(qwen3_model_name, dflash_output_dir):
     first_loss = float(logs[0]["loss"])
     final_loss = float(logs[-1]["loss"])
     assert final_loss < first_loss, f"Loss did not decrease: {first_loss:.3f} -> {final_loss:.3f}"
-    # Sanity: final loss should be reasonable (baseline: ~1.1 on L40)
+    # Sanity: final loss should be reasonable (baseline: ~1.8 at step 500 on L40)
     assert final_loss < 3.0, f"Final loss {final_loss:.3f} too high (expected < 3.0)"
 
 
@@ -118,7 +122,7 @@ def test_dflash_resume(qwen3_model_name, dflash_output_dir):
     overrides = [
         f"model.model_name_or_path={qwen3_model_name}",
         f"training.output_dir={output_dir}",
-        "training.num_train_epochs=4",
+        f"training.max_steps={_TRAIN_STEPS + _RESUME_STEPS}",
         "training.save_steps=5000",
         *_DFLASH_OVERRIDES,
     ]
@@ -128,6 +132,11 @@ def test_dflash_resume(qwen3_model_name, dflash_output_dir):
         "speculative_decoding",
     )
 
+    with open(os.path.join(output_dir, "trainer_state.json")) as f:
+        state = json.load(f)
+    assert state["global_step"] == _TRAIN_STEPS + _RESUME_STEPS, state["global_step"]
+    assert any(h["step"] > _TRAIN_STEPS for h in state["log_history"]), "no log entry after resume"
+
 
 def test_dflash_export(dflash_output_dir):
     """Export DFlash checkpoint to deployment format."""

From c4872a810a15b1323cc521dc94c337c03945de09 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:44:44 +0000
Subject: [PATCH 13/34] Skip the Megatron state reset for megatron_bridge tests
 that use no GPU

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>
---
 tests/examples/megatron_bridge/conftest.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
index 1093fb7b281..17043f6aebc 100644
--- a/tests/examples/megatron_bridge/conftest.py
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -40,18 +40,22 @@ def _fast_example_runner():
 
 
 @pytest.fixture(autouse=True)
-def _isolate_megatron_global_state():
+def _isolate_megatron_global_state(request):
     """Reset shared state around every test so a failure cannot cascade into the next one.
 
     In-process steps share the interpreter. Besides Megatron's singletons, Transformer-Engine
     records its chosen attention backend in ``NVTE_*``, which failed a Mamba hybrid that ran after
     an attention model -- so the environment is restored wholesale rather than by naming variables.
     """
+    # Only tests that take ``num_gpus`` run example steps; the gc-heavy reset is wasted on the rest.
+    reset_state = "num_gpus" in request.fixturenames
     env_before = os.environ.copy()
-    reset_megatron_global_state()
+    if reset_state:
+        reset_megatron_global_state()
     try:
         yield
     finally:
-        reset_megatron_global_state()
+        if reset_state:
+            reset_megatron_global_state()
         os.environ.clear()
         os.environ.update(env_before)

From cb1fc123f593af96f85ec85e4cc768b3e2c6d1dd Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:44:58 +0000
Subject: [PATCH 14/34] Use the tiny local tokenizer in six preprocess_data
 tests

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>
---
 .../plugins/test_megatron_preprocess_data.py  | 26 ++++++++++---------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
index 5dd0815080c..be96e50e037 100644
--- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
+++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
@@ -49,7 +49,7 @@ def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
+def test_megatron_preprocess_data_with_jsonl_path(tmp_path, tiny_tokenizer_path):
     input_jsonl = tmp_path / "minipile.jsonl"
     input_jsonl.write_text(
         "".join(json.dumps({"text": f"Sample document {index}."}) + "\n" for index in range(4)),
@@ -67,7 +67,7 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     prefixes = megatron_preprocess_data(
         jsonl_paths=input_jsonl,
         output_dir=tmp_path,
-        tokenizer_name_or_path="gpt2",
+        tokenizer_name_or_path=tiny_tokenizer_path,
         json_keys=["text"],
         workers=1,
     )
@@ -81,7 +81,7 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     assert Path(prefixes[0] + ".idx").stat().st_size > 0, "Index file should not be empty"
 
 
-def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
+def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path, tiny_tokenizer_path):
     input_path = tmp_path / "data.jsonl"
     input_path.write_text(
         "".join(json.dumps({"text": f"Document {index} " * 100}) + "\n" for index in range(10)),
@@ -91,7 +91,7 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
     common_args = {
         "jsonl_paths": input_path,
         "output_dir": tmp_path,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "workers": 2,
     }
@@ -115,13 +115,13 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
+def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path, tiny_tokenizer_path):
     common_args = {
         "hf_dataset": "nanotron/minipile_100_samples",
         "hf_split": "train",
         "hf_max_samples_per_split": 100,
         "hf_streaming": True,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "workers": 2,
     }
@@ -147,7 +147,9 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path):
+def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(
+    tmp_path, tiny_tokenizer_path
+):
     args = {
         "hf_dataset": "nanotron/minipile_100_samples",
         "hf_split": "train",
@@ -155,7 +157,7 @@ def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_pa
         "hf_streaming": True,
         "max_tokens": 100,
         "output_dir": tmp_path,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "max_sequence_length": 16,
         "workers": 2,
@@ -253,7 +255,7 @@ def test_megatron_preprocess_data_reasoning_content(tmp_path):
     )
 
 
-def test_megatron_preprocess_data_with_gzip_input(tmp_path):
+def test_megatron_preprocess_data_with_gzip_input(tmp_path, tiny_tokenizer_path):
     gz_path = tmp_path / "data.jsonl.gz"
     with gzip.open(gz_path, "wt", encoding="utf-8") as f:
         for i in range(10):
@@ -262,7 +264,7 @@ def test_megatron_preprocess_data_with_gzip_input(tmp_path):
     prefixes = megatron_preprocess_data(
         jsonl_paths=gz_path,
         output_dir=tmp_path,
-        tokenizer_name_or_path="gpt2",
+        tokenizer_name_or_path=tiny_tokenizer_path,
         json_keys=["text"],
         workers=1,
         strip_newlines=True,
@@ -324,7 +326,7 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_streaming_warning(tmp_path):
+def test_megatron_preprocess_data_hf_streaming_warning(tmp_path, tiny_tokenizer_path):
     # hf_streaming without hf_max_samples_per_split should warn and fall back to non-streaming
     with pytest.warns(UserWarning, match="hf_streaming"):
         megatron_preprocess_data(
@@ -332,7 +334,7 @@ def test_megatron_preprocess_data_hf_streaming_warning(tmp_path):
             hf_split="train",
             hf_streaming=True,
             output_dir=tmp_path,
-            tokenizer_name_or_path="gpt2",
+            tokenizer_name_or_path=tiny_tokenizer_path,
             json_keys=["text"],
             workers=1,
         )

From 9536e27fee0cf77c1e394252d12da79fdb47fe65 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 9 Oct 2026 23:59:08 +0000
Subject: [PATCH 15/34] Run the CPU-only ONNX quantization tests in the unit
 lane

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>
---
 tests/{gpu => unit}/onnx/quantization/test_dla_qdq_rules.py       | 0
 .../{gpu => unit}/onnx/quantization/test_dla_scalar_constants.py  | 0
 .../onnx/quantization/test_quantization_initializer_names.py      | 0
 3 files changed, 0 insertions(+), 0 deletions(-)
 rename tests/{gpu => unit}/onnx/quantization/test_dla_qdq_rules.py (100%)
 rename tests/{gpu => unit}/onnx/quantization/test_dla_scalar_constants.py (100%)
 rename tests/{gpu => unit}/onnx/quantization/test_quantization_initializer_names.py (100%)

diff --git a/tests/gpu/onnx/quantization/test_dla_qdq_rules.py b/tests/unit/onnx/quantization/test_dla_qdq_rules.py
similarity index 100%
rename from tests/gpu/onnx/quantization/test_dla_qdq_rules.py
rename to tests/unit/onnx/quantization/test_dla_qdq_rules.py
diff --git a/tests/gpu/onnx/quantization/test_dla_scalar_constants.py b/tests/unit/onnx/quantization/test_dla_scalar_constants.py
similarity index 100%
rename from tests/gpu/onnx/quantization/test_dla_scalar_constants.py
rename to tests/unit/onnx/quantization/test_dla_scalar_constants.py
diff --git a/tests/gpu/onnx/quantization/test_quantization_initializer_names.py b/tests/unit/onnx/quantization/test_quantization_initializer_names.py
similarity index 100%
rename from tests/gpu/onnx/quantization/test_quantization_initializer_names.py
rename to tests/unit/onnx/quantization/test_quantization_initializer_names.py

From c5af12e73f24f40b6f56855101c3f647b871be65 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:00:15 +0000
Subject: [PATCH 16/34] Calibrate test_concat_elim on CPU with a small model

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>
---
 tests/_test_utils/onnx/lib_test_models.py     | 22 +++++++++----------
 .../gpu/onnx/quantization/test_concat_elim.py | 17 ++++++++++----
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/tests/_test_utils/onnx/lib_test_models.py b/tests/_test_utils/onnx/lib_test_models.py
index 57a0621c5da..12dba85b20f 100644
--- a/tests/_test_utils/onnx/lib_test_models.py
+++ b/tests/_test_utils/onnx/lib_test_models.py
@@ -260,12 +260,12 @@ def find_init(onnx_model: onnx.ModelProto, init_name: str) -> np.ndarray:
 #           conv_4
 #             |
 #
-def build_conv_concat_model():
+def build_conv_concat_model(channels=128, height=240, width=64):
     # Define your model inputs and outputs
     input_names = ["input_0"]
     output_names = ["conv4_conv/Conv2D:0"]
-    input_shapes = [(1, 128, 240, 64)]
-    output_shapes = [(1, 256, 240, 64)]
+    input_shapes = [(1, channels, height, width)]
+    output_shapes = [(1, 2 * channels, height, width)]
 
     inputs = [
         helper.make_tensor_value_info(input_name, onnx.TensorProto.FLOAT, input_shape)
@@ -340,26 +340,26 @@ def build_conv_concat_model():
         helper.make_tensor(
             name="ReadVariableOp:0",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:1",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:2",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:3",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(256, 384, 1, 1),
-            vals=np.random.uniform(low=0.5, high=1.0, size=256 * 384 * 1 * 1),
+            dims=(2 * channels, 3 * channels, 1, 1),
+            vals=np.random.uniform(low=0.5, high=1.0, size=2 * channels * 3 * channels),
         ),
     ]
 
diff --git a/tests/gpu/onnx/quantization/test_concat_elim.py b/tests/gpu/onnx/quantization/test_concat_elim.py
index 68179051909..de9feafe7e7 100644
--- a/tests/gpu/onnx/quantization/test_concat_elim.py
+++ b/tests/gpu/onnx/quantization/test_concat_elim.py
@@ -23,10 +23,19 @@
 
 from modelopt.onnx.quantization.quantize import quantize
 
+# Keep every Conv above the 16-channel limit for INT8/FP8 quantization (graph_selection.py) so the
+# same nodes are quantized as with the default 128-channel model, but calibrate a small input.
+SMALL_MODEL_SHAPE = {"channels": 32, "height": 16, "width": 16}
+
 
 def _check_concat_qdq_status(onnx_path, quantize_mode):
-    # Quantize the input model
-    quantize(onnx_path, quantize_mode=quantize_mode, passes="concat_elimination")
+    # Quantize the input model. Only the graph rewrite is checked, so calibrate on CPU.
+    quantize(
+        onnx_path,
+        quantize_mode=quantize_mode,
+        passes="concat_elimination",
+        calibration_eps=["cpu"],
+    )
 
     # Output model should be produced in the same tmpdir
     output_onnx_path = onnx_path.replace(".onnx", ".quant.onnx")
@@ -61,7 +70,7 @@ def _check_concat_qdq_status(onnx_path, quantize_mode):
 
 def test_concat_elim_int8(tmp_path):
     # Copy the test model to the tmp_path
-    model = build_conv_concat_model()
+    model = build_conv_concat_model(**SMALL_MODEL_SHAPE)
     this_function_name = sys._getframe().f_code.co_name
     onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx")
     onnx.save(model, onnx_path)
@@ -70,7 +79,7 @@ def test_concat_elim_int8(tmp_path):
 
 def test_concat_elim_fp8(tmp_path):
     # Copy the test model to the tmp_path
-    model = build_conv_concat_model()
+    model = build_conv_concat_model(**SMALL_MODEL_SHAPE)
     this_function_name = sys._getframe().f_code.co_name
     onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx")
     onnx.save(model, onnx_path)

From 1662f965d6cff44fca396b4691aba9c50fb88073 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:01:17 +0000
Subject: [PATCH 17/34] Benchmark a small model in the autotune benchmark tests

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>
---
 .../onnx/quantization/autotune/models.py       |  6 +++---
 .../quantization/autotune/test_benchmark.py    | 18 ++++++++++++++----
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/tests/_test_utils/onnx/quantization/autotune/models.py b/tests/_test_utils/onnx/quantization/autotune/models.py
index 84a8b4ab82e..89721661b4a 100644
--- a/tests/_test_utils/onnx/quantization/autotune/models.py
+++ b/tests/_test_utils/onnx/quantization/autotune/models.py
@@ -25,13 +25,13 @@
 from onnx import helper
 
 
-def _create_simple_conv_onnx_model():
+def _create_simple_conv_onnx_model(batch_size=64, spatial_size=224):
     """Build ONNX model: Input -> Conv -> Relu -> Output (minimal for autotuner tests)."""
     input_tensor = helper.make_tensor_value_info(
-        "input", onnx.TensorProto.FLOAT, [64, 32, 224, 224]
+        "input", onnx.TensorProto.FLOAT, [batch_size, 32, spatial_size, spatial_size]
     )
     output_tensor = helper.make_tensor_value_info(
-        "output", onnx.TensorProto.FLOAT, [64, 64, 224, 224]
+        "output", onnx.TensorProto.FLOAT, [batch_size, 64, spatial_size, spatial_size]
     )
     conv_node = helper.make_node(
         "Conv",
diff --git a/tests/gpu/onnx/quantization/autotune/test_benchmark.py b/tests/gpu/onnx/quantization/autotune/test_benchmark.py
index 012bf495ab6..47311e25219 100644
--- a/tests/gpu/onnx/quantization/autotune/test_benchmark.py
+++ b/tests/gpu/onnx/quantization/autotune/test_benchmark.py
@@ -46,8 +46,8 @@
 
 @pytest.fixture
 def simple_conv_model_bytes():
-    """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils)."""
-    model = _create_simple_conv_onnx_model()
+    """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils), small input."""
+    model = _create_simple_conv_onnx_model(batch_size=1, spatial_size=32)
     return model.SerializeToString()
 
 
@@ -431,7 +431,12 @@ def test_run_with_path(self, simple_conv_model_path):
         with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f:
             cache_path = f.name
         try:
-            benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2)
+            benchmark = TrtExecBenchmark(
+                timing_cache_file=cache_path,
+                warmup_runs=1,
+                timing_runs=2,
+                trtexec_args=["--duration=0"],
+            )
             latency_ms = benchmark.run(simple_conv_model_path)
             assert isinstance(latency_ms, float)
             assert latency_ms > 0
@@ -445,7 +450,12 @@ def test_run_with_bytes(self, simple_conv_model_bytes):
         with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f:
             cache_path = f.name
         try:
-            benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2)
+            benchmark = TrtExecBenchmark(
+                timing_cache_file=cache_path,
+                warmup_runs=1,
+                timing_runs=2,
+                trtexec_args=["--duration=0"],
+            )
             latency_ms = benchmark.run(simple_conv_model_bytes)
             assert isinstance(latency_ms, float)
             assert latency_ms > 0

From aea663a8acb3fd99f43b6874acd83ad8af07920d Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:01:52 +0000
Subject: [PATCH 18/34] Skip trtexec's 3-second measurement soak in the
 autotune workflow test

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>
---
 tests/gpu/onnx/quantization/autotune/test_workflow.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/gpu/onnx/quantization/autotune/test_workflow.py b/tests/gpu/onnx/quantization/autotune/test_workflow.py
index eebc42f6c33..b323402f1eb 100644
--- a/tests/gpu/onnx/quantization/autotune/test_workflow.py
+++ b/tests/gpu/onnx/quantization/autotune/test_workflow.py
@@ -59,7 +59,10 @@ def test_export_quantized_model(use_trtexec, simple_conv_model):
     output_path = output_dir + ".quant.onnx"
 
     try:
-        init_benchmark_instance(use_trtexec=use_trtexec, timing_runs=100)
+        # trtexec measures for at least 3 s by default; --duration=0 stops after timing_runs iterations
+        init_benchmark_instance(
+            use_trtexec=use_trtexec, timing_runs=100, trtexec_args=["--duration=0"]
+        )
         autotuner = region_pattern_autotuning_workflow(baseline_model_path, Path(output_dir))
 
         # Export model with Q/DQ insertion

From 7aa6a2173e410711e1a8b329a06a280fcafee029 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:02:41 +0000
Subject: [PATCH 19/34] Shrink the puzzletron test data and fold a duplicate
 NAS test into another

- 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>
---
 tests/_test_utils/torch/puzzletron/utils.py   |  6 +-
 .../Qwen2.5-7B-Instruct.yaml                  |  2 +-
 .../configs/Qwen/Qwen3-8B/Qwen3-8B.yaml       |  2 +-
 .../Qwen3-VL-30B-A3B-Instruct.yaml            |  2 +-
 .../Llama-3.1-8B-Instruct.yaml                |  2 +-
 .../Llama-3.2-3B-Instruct.yaml                |  2 +-
 .../Mistral-Small-24B-Instruct-2501.yaml      |  2 +-
 ...DIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml |  2 +-
 .../NVIDIA-Nemotron-Nano-12B-v2.yaml          |  2 +-
 .../openai/gpt-oss-20b/gpt-oss-20b.yaml       |  2 +-
 .../gpu/torch/puzzletron/test_nas_convert.py  | 56 -------------------
 tests/gpu/torch/puzzletron/test_nas_search.py |  9 ++-
 12 files changed, 20 insertions(+), 69 deletions(-)

diff --git a/tests/_test_utils/torch/puzzletron/utils.py b/tests/_test_utils/torch/puzzletron/utils.py
index 2ce44858ef5..525b597518d 100644
--- a/tests/_test_utils/torch/puzzletron/utils.py
+++ b/tests/_test_utils/torch/puzzletron/utils.py
@@ -270,9 +270,9 @@ def save_dummy_dataset(dataset_path: Path | str):
         },
     ]
 
-    # Prepare train and val splits with sample repeated, 2500 samples are for
-    # 128 samples with block-size 8192 and LLama3 tokenizer
-    data = [{"conversation": sample}] * 2500
+    # Prepare train and val splits with the sample repeated. With the tiny tokenizer a sample is
+    # 1424 tokens, so 400 samples give 278 blocks of 2048 tokens; the pruning scores read at most 100.
+    data = [{"conversation": sample}] * 400
 
     # For train-val splits
     data_dict = DatasetDict({"train": Dataset.from_list(data), "valid": Dataset.from_list(data)})
diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml
index 2843f0b97ac..6691a65d9d7 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml
index cd82a47271a..59e69520a4b 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml
index 00b21ea9798..77567a777f9 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml
index 8e2e0786b33..880f630d83b 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml
index 78cb6bd73c4..ab45abe41f9 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml b/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml
index e042c4bb627..a6593200354 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml
index ab2b09e6794..33a8318215d 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml
@@ -20,7 +20,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml
index 906b7338d8c..05b22051392 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml
@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml b/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml
index 60d6568df60..d57ff295796 100644
--- a/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml
+++ b/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml
@@ -20,7 +20,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths
diff --git a/tests/gpu/torch/puzzletron/test_nas_convert.py b/tests/gpu/torch/puzzletron/test_nas_convert.py
index 57239f1c0e4..7ddb0584fc7 100644
--- a/tests/gpu/torch/puzzletron/test_nas_convert.py
+++ b/tests/gpu/torch/puzzletron/test_nas_convert.py
@@ -27,62 +27,6 @@
 import modelopt.torch.utils.distributed as dist
 
 
-def test_nas_convert_ffn_pruning(project_root_path: Path, tmp_path: Path):
-    spawn_multiprocess_job(
-        size=torch.cuda.device_count(),
-        job=partial(_test_nas_convert_ffn_pruning_multiprocess_job, project_root_path, tmp_path),
-        backend="nccl",
-    )
-
-
-def _test_nas_convert_ffn_pruning_multiprocess_job(
-    project_root_path: Path, tmp_path: Path, rank: int, size: int
-):
-    dist.setup(timeout=timedelta(minutes=10))
-    # Setup the test model and data.
-    puzzle_dir, llama_checkpoint_path, dataset_path = setup_test_model_and_data(
-        tmp_path, rank, "meta-llama/Llama-3.1-8B-Instruct"
-    )
-    hydra_config_dir = project_root_path / "tests/gpu/torch/puzzletron/resources/configs"
-    hydra_config_name = "meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct"
-
-    #
-    # Run the mnt.convert() step
-    #
-    input_model = mtpz.puzzletron_nas_plugin.PuzzletronModel()
-    mtn.convert(
-        input_model,
-        mode=[
-            (
-                "puzzletron",
-                {
-                    "puzzle_dir": str(puzzle_dir),
-                    "input_model_path": str(llama_checkpoint_path),
-                    "hydra_config_dir": str(hydra_config_dir),
-                    "hydra_config_name": hydra_config_name,
-                    "dataset_path": str(dataset_path),
-                },
-            )
-        ],
-    )
-
-    #
-    # Check assertions
-    #
-    if rank == 0:
-        # assertions for the score_pruning_activations step
-        rank = int(os.environ["RANK"])
-        rank_filepath = (
-            f"pruning/pruning_scores/ffn_iterative/100samples_diverse_mini/rank_{rank}.pth"
-        )
-        assert (puzzle_dir / rank_filepath).is_file()
-
-        # assertions for the pruning_ckpts step
-        assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists()
-
-    dist.cleanup()
-
-
 def test_nas_convert_attn_pruning(project_root_path: Path, tmp_path: Path):
     spawn_multiprocess_job(
         size=torch.cuda.device_count(),
diff --git a/tests/gpu/torch/puzzletron/test_nas_search.py b/tests/gpu/torch/puzzletron/test_nas_search.py
index ae0dff47be5..4cbd95ca2f9 100644
--- a/tests/gpu/torch/puzzletron/test_nas_search.py
+++ b/tests/gpu/torch/puzzletron/test_nas_search.py
@@ -76,9 +76,16 @@ def _test_nas_search_multiprocess_job(
     )
 
     #
-    # Check assertions for mtn.search() step
+    # Check assertions for the mtn.convert() and mtn.search() steps
     #
     if rank == 0:
+        # assertions for the score_pruning_activations step of mtn.convert()
+        rank_filepath = "pruning/pruning_scores/ffn_iterative/100samples_diverse_mini/rank_0.pth"
+        assert (puzzle_dir / rank_filepath).is_file()
+
+        # assertions for the pruning_ckpts step of mtn.convert()
+        assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists()
+
         # assertions for the build_library_and_stats step
         assert (puzzle_dir / "replacement_library.json").is_file()
         assert (puzzle_dir / "subblock_stats.json").is_file()

From b3bdaf0f4da78f3954c957b100a11d14562f1208 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:03:30 +0000
Subject: [PATCH 20/34] Calibrate the unified HF export test on local text,
 with no previews 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>
---
 tests/_test_utils/torch/calib_data.py         | 41 +++++++++++++++++++
 ...unified_hf_export_and_check_safetensors.py | 14 ++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/_test_utils/torch/calib_data.py

diff --git a/tests/_test_utils/torch/calib_data.py b/tests/_test_utils/torch/calib_data.py
new file mode 100644
index 00000000000..03622986d82
--- /dev/null
+++ b/tests/_test_utils/torch/calib_data.py
@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Local calibration text for tests that only need some text, not a particular Hub dataset."""
+
+import json
+import random
+from pathlib import Path
+
+__all__ = ["write_calib_jsonl"]
+
+_TEXT = (
+    "the of and to in is that for it as was with be by on not he this are or his from at which "
+    "but have an had they you were their one all we can her has there been if more when will "
+    "would who so no model layer weight scale range value sample batch token cache quantize"
+)
+_WORDS = _TEXT.split()
+
+
+def write_calib_jsonl(path: Path | str, num_samples: int = 64, words_per_sample: int = 150) -> str:
+    """Write ``num_samples`` random-word texts to ``path`` as ``{"text": ...}`` JSON lines.
+
+    Every text is long enough to fill the default 512-token calibration window with a character
+    level tokenizer. The texts differ from each other and are the same on every call.
+    """
+    rng = random.Random(0)
+    texts = [" ".join(rng.choices(_WORDS, k=words_per_sample)) for _ in range(num_samples)]
+    Path(path).write_text("".join(json.dumps({"text": t}) + "\n" for t in texts), encoding="utf-8")
+    return str(path)
diff --git a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py
index b45943e6e5e..7a67c026b3b 100644
--- a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py
+++ b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py
@@ -17,6 +17,7 @@
 import pytest
 import torch
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
+from _test_utils.torch.calib_data import write_calib_jsonl
 from _test_utils.torch.transformers_models import (
     create_tiny_gpt_oss_dir,
     create_tiny_llama_dir,
@@ -26,6 +27,12 @@
 from safetensors import safe_open
 
 
+@pytest.fixture(scope="module")
+def calib_jsonl(tmp_path_factory):
+    """Local calibration texts; the test checks the exported structure, so it needs no Hub dataset."""
+    return write_calib_jsonl(tmp_path_factory.mktemp("calib") / "calib.jsonl")
+
+
 # Here we map each qformat -> the suffix we expect in the generated safetensors directory
 @pytest.mark.parametrize(
     (
@@ -55,6 +62,7 @@
 )
 def test_unified_hf_export_and_check_safetensors(
     tmp_path,
+    calib_jsonl,
     qformat,
     expected_suffix,
     fuse_input_scale,
@@ -103,11 +111,15 @@ def test_unified_hf_export_and_check_safetensors(
         pyt_ckpt_path=tiny_model_dir,
         qformat=qformat,
         export_path=output_dir,
-        dataset="cnn_dailymail",
+        dataset=calib_jsonl,
         # This test only checks the exported safetensors structure (not accuracy), so a
         # small calibration set is enough. Avoids the 1024-sample default on a toy model.
         calib_size=64,
+        # One batch holds the whole set, so hf_ptq does not probe for the largest batch.
+        batch_size=64,
     )
+    # The generate() previews before and after quantization do not affect the export.
+    cmd_parts.append("--skip_generate")
 
     # Run the command
     # current transformers T5 model seem to be incompatible with multiple GPUs

From 72c16fe22b55d5f5d12d33628be74fdfb340147a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:03:50 +0000
Subject: [PATCH 21/34] Calibrate the GPTQ end-to-end test on local text

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>
---
 tests/gpu/torch/quantization/test_gptq.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index 496db08578a..abe3386064a 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+from _test_utils.torch.calib_data import write_calib_jsonl
 from _test_utils.torch.transformers_models import get_tiny_llama
 from conftest import requires_triton
 
@@ -130,14 +131,14 @@ def test_gptq_export_roundtrip():
 @pytest.mark.parametrize(
     "quant_cfg", [mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG, mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG]
 )
-def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer):
+def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer, tmp_path):
     model = get_tiny_llama(vocab_size=tiny_tokenizer.vocab_size).to("cuda")
     model.eval()
 
     quant_cfg = copy.deepcopy(quant_cfg)
     quant_cfg["algorithm"] = {"method": "gptq", "layerwise": {"enable": True}}
     calib_dataloader = get_dataset_dataloader(
-        dataset_name="cnn_dailymail",
+        dataset_name=write_calib_jsonl(tmp_path / "calib.jsonl", num_samples=8),
         tokenizer=tiny_tokenizer,
         batch_size=2,
         num_samples=8,

From 35b7eaecf01d535bb3dd0f454b2d2204af66539b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:04:46 +0000
Subject: [PATCH 22/34] Evaluate one MMLU subtask in test_lm_eval_hf

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>
---
 tests/examples/llm_eval/test_llm_eval.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py
index d65f497e320..ec28ab28aa9 100644
--- a/tests/examples/llm_eval/test_llm_eval.py
+++ b/tests/examples/llm_eval/test_llm_eval.py
@@ -32,9 +32,10 @@ def test_lm_eval_hf(tmp_path):
         ["python", "lm_eval_hf.py"],
         model="hf",
         model_args=f"pretrained={model_dir}",
-        tasks="mmlu",
+        # One subtask takes the same harness and gate path as the 57-subtask mmlu group, which
+        # spends minutes loading datasets on a model whose context is only 32 tokens
+        tasks="mmlu_abstract_algebra",
         num_fewshot=5,
-        limit=0.1,
         batch_size=8,
         # Exercise the accuracy gate: reads the results file and enforces the bound
         output_path=str(tmp_path / "results"),

From 43eb82e832513fceca7554fa8fbfbd263df8649a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:04:52 +0000
Subject: [PATCH 23/34] Use one MMLU prompt per subject in test_qwen3_eval_fp8

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>
---
 tests/examples/llm_eval/test_llm_eval.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py
index ec28ab28aa9..c1ba2c29e95 100644
--- a/tests/examples/llm_eval/test_llm_eval.py
+++ b/tests/examples/llm_eval/test_llm_eval.py
@@ -63,9 +63,9 @@ def test_qwen3_eval_fp8(tmp_path):
             simple_eval_tasks="humaneval",
             lm_eval_limit=16,
             simple_eval_limit=16,
-            # MMLU has no inherent sample cap and otherwise evaluates the full ~14k-question
-            # test set; limit it like lm_eval/simple_eval to keep this a fast smoke test.
-            mmlu_limit=16,
+            # mmlu.py applies the limit per subject (57 of them), one TRT-LLM call per prompt:
+            # 1 still exercises every subject, 16 would be 912 sequential calls.
+            mmlu_limit=1,
             output=128,  # Cap generation length: gsm8k/humaneval otherwise generate up to 1024 tokens/sample
             batch=8,
         )

From 987ad4c77280e4e6ce9778946b541c48212ab318 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:05:20 +0000
Subject: [PATCH 24/34] Quantize once per module in the llm_qat QAT and QAD
 tests

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>
---
 tests/examples/llm_qat/test_llm_qat.py | 85 +++++++++++++-------------
 1 file changed, 44 insertions(+), 41 deletions(-)

diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index a7b610a807c..bec89bd1b97 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/tests/examples/llm_qat/test_llm_qat.py
@@ -19,6 +19,7 @@
 import pytest
 import torch
 from _test_utils.examples.run_command import run_example_command
+from _test_utils.fs_utils import assert_unmodified_tree
 from safetensors.torch import load_file
 
 # Mapping from backend name to accelerate config file
@@ -102,6 +103,38 @@ def _run_export(ckpt_dir: str, export_dir: str):
     )
 
 
+@pytest.fixture(scope="module")
+def dataset_cache_dir(tmp_path_factory):
+    """Tokenized-dataset cache for every quantize/train run below.
+
+    The cache key covers the tokenizer, max length and blend (not the command), so the dataset is
+    built once for quantize.py (8192) and once for train.py (128) instead of once per command.
+    """
+    return str(tmp_path_factory.mktemp("dataset_cache"))
+
+
+@pytest.fixture(scope="module")
+def nvfp4_ptq_dir(tiny_qwen3_path, tmp_path_factory, dataset_cache_dir):
+    """NVFP4 PTQ checkpoint shared by the QAT and QAD tests, which only read it.
+
+    quantize.py ignores the QAD-only config keys, so qat_nvfp4.yaml and qad_nvfp4.yaml give the
+    same arguments and one run serves both.
+    """
+    ptq_dir = tmp_path_factory.mktemp("ptq")
+    _run_quantize(
+        "configs/train/qat_nvfp4.yaml",
+        [
+            "--model_name_or_path", tiny_qwen3_path,
+            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
+            "--calib_size", "64",
+            "--output_dir", str(ptq_dir),
+        ],
+        cache_dir=dataset_cache_dir,
+    )
+    with assert_unmodified_tree(ptq_dir) as path:
+        yield str(path)
+
+
 def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path):
     """Test dataset_utils.py standalone CLI pre-tokenization."""
     cache_dir = tmp_path / "dataset_cache"
@@ -122,33 +155,18 @@ def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path):
     "deepspeed",
     "ddp",
 ])
-def test_qwen3_qat_nvfp4(tiny_qwen3_path, tmp_path, backend):
-    ptq_output_dir = tmp_path / "ptq"
+def test_qwen3_qat_nvfp4(nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend):
     qat_output_dir = tmp_path / "qat"
-    cache_dir = str(tmp_path / "dataset_cache")
-
-    # Step 1: Quantize
-    _run_quantize(
-        "configs/train/qat_nvfp4.yaml",
-        [
-            "--model_name_or_path", tiny_qwen3_path,
-            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
-            "--calib_size", "64",
-            "--output_dir", str(ptq_output_dir),
-        ],
-        cache_dir=cache_dir,
-    )
 
-    # Step 2: QAT
     _run_train(
         "configs/train/qat_nvfp4.yaml",
         [
-            "--model_name_or_path", str(ptq_output_dir),
+            "--model_name_or_path", nvfp4_ptq_dir,
             "--do_train", "True",
             "--output_dir", str(qat_output_dir),
         ],
         backend=backend,
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
 @pytest.mark.skip(reason="FSDP2 LoRA checkpoint save omits adapter_model.safetensors")
@@ -211,43 +229,28 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     "fsdp2",
     "deepspeed",
 ])
-def test_qwen3_qad_nvfp4(tiny_qwen3_path, tmp_path, backend):
-    ptq_output_dir = tmp_path / "ptq"
+def test_qwen3_qad_nvfp4(tiny_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend):
     qad_output_dir = tmp_path / "qad"
-    cache_dir = str(tmp_path / "dataset_cache")
-
-    # Step 1: Quantize student
-    _run_quantize(
-        "configs/train/qad_nvfp4.yaml",
-        [
-            "--model_name_or_path", tiny_qwen3_path,
-            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
-            "--calib_size", "64",
-            "--output_dir", str(ptq_output_dir),
-        ],
-        cache_dir=cache_dir,
-    )
 
-    # Step 2: QAD (quantization-aware distillation)
+    # QAD (quantization-aware distillation) of the shared quantized student
     _run_train(
         "configs/train/qad_nvfp4.yaml",
         [
-            "--model_name_or_path", str(ptq_output_dir),
+            "--model_name_or_path", nvfp4_ptq_dir,
             "--do_train", "True",
             "--output_dir", str(qad_output_dir),
             "--distill", "True",
             "--teacher_model", tiny_qwen3_path,
         ],
         backend=backend,
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
 
-def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
+def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path):
     ptq_output_dir = tmp_path / "ptq"
-    cache_dir = str(tmp_path / "dataset_cache")
 
-    # Step 1: Quantize with compression for QLoRA
+    # Step 1: Quantize with compression for QLoRA (its own checkpoint: the others are not compressed)
     _run_quantize(
         "configs/train/qlora_nvfp4.yaml",
         [
@@ -257,7 +260,7 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
             "--compress", "True",
             "--output_dir", str(ptq_output_dir),
         ],
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
     # Step 2: QLoRA training
@@ -271,7 +274,7 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
             "--output_dir", str(qlora_output_dir),
         ],
         backend="ddp",
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
     # Step 3: Export the QLoRA checkpoint for deployment

From 3e60ef15eac14fa90575f4b421b569d3dc8ec7b3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:05:31 +0000
Subject: [PATCH 25/34] Skip Triton GEMM autotune in the 1-GPU EAGLE3 compile
 smoke test

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>
---
 tests/examples/speculative_decoding/test_eagle.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py
index 8d284d4ba1b..30f88ec985b 100644
--- a/tests/examples/speculative_decoding/test_eagle.py
+++ b/tests/examples/speculative_decoding/test_eagle.py
@@ -147,9 +147,17 @@ def test_llama_eagle3(tiny_llama_path,
         *_TINY_EAGLE_ARCH,
     ]
 
+    env = None
+    if cp_size == 1 and not mix_hidden_states:
+        # The 1-GPU PR smoke compiles with mode="max-autotune"; benchmarking Triton GEMM templates
+        # is a large share of that on a tiny model and the test only needs the compiled path to
+        # run. [2-False] (2 GPUs, nightly) keeps the full autotune.
+        env = {**os.environ, "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS": "ATEN"}
+
     run_example_command(
         ["./launch_train.sh", "--config", EAGLE3_YAML, *overrides],
         "speculative_decoding",
+        env=env,
     )
 
 

From e53b909cf0de2f3d5577521b9492da579d803dcf Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:06:01 +0000
Subject: [PATCH 26/34] Drop the duplicate tiny-llama offline EAGLE3 training
 test

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>
---
 .../examples/speculative_decoding/test_eagle.py  | 16 ++++++++--------
 .../test_eagle_offline_ptq.py                    |  3 ++-
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py
index 30f88ec985b..7325fe72e84 100644
--- a/tests/examples/speculative_decoding/test_eagle.py
+++ b/tests/examples/speculative_decoding/test_eagle.py
@@ -230,7 +230,6 @@ def test_convert_to_vllm_ckpt(tiny_llama_path, eagle_output_dir):
 @pytest.mark.parametrize(
     ("model_source", "use_fake_base"),
     [
-        (None, False),  # tiny_llama (from fixture), no FakeBase
         pytest.param("moonshotai/Kimi-K2.5", True),  # remote HF repo, FakeBaseModel
         pytest.param(  # remote HF repo, no FakeBaseModel
             "moonshotai/Kimi-K2-Thinking",
@@ -243,22 +242,23 @@ def test_convert_to_vllm_ckpt(tiny_llama_path, eagle_output_dir):
             marks=pytest.mark.manual(reason="skip redundand test, too slow"),
         ),
     ],
-    ids=["tinyllama", "kimi-k2.5", "kimi-k2-thinking", "minimax-m2.5"],
+    ids=["kimi-k2.5", "kimi-k2-thinking", "minimax-m2.5"],
 )
 def test_offline_eagle3_training(
-    tiny_llama_path,
     tiny_daring_anteater_path,
     tmp_path,
     eagle_output_dir,
     model_source,
     use_fake_base,
 ):
-    """Test Eagle3 training with pre-computed hidden states (offline mode / FakeBaseModel)."""
-    model_path = tiny_llama_path if model_source is None else model_source
-    model_id = "tinyllama" if model_source is None else model_source.split("/")[-1]
+    """Test Eagle3 training with pre-computed hidden states on a FakeBaseModel.
+
+    Offline training on a real base model is covered by test_eagle_offline_ptq.py.
+    """
+    model_id = model_source.split("/")[-1]
     output_subdir = eagle_output_dir / f"eagle-{model_id}-offline"
 
-    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    cfg = AutoConfig.from_pretrained(model_source, trust_remote_code=True)
     if hasattr(cfg, "text_config"):  # vlm: get text_config
         cfg = cfg.text_config
 
@@ -270,7 +270,7 @@ def test_offline_eagle3_training(
     )
 
     overrides = [
-        f"model.model_name_or_path={model_path}",
+        f"model.model_name_or_path={model_source}",
         "model.trust_remote_code=true",
         f"model.use_fake_base_for_offline={str(use_fake_base).lower()}",
         f"data.data_path={tiny_daring_anteater_path}",
diff --git a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py b/tests/examples/speculative_decoding/test_eagle_offline_ptq.py
index 5fca5a3d71c..3bd1be71dba 100644
--- a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py
+++ b/tests/examples/speculative_decoding/test_eagle_offline_ptq.py
@@ -97,7 +97,8 @@ def test_offline_eagle_training(tiny_llama_path, tiny_daring_anteater_path, offl
         f"training.output_dir={output_dir}",
         "training.num_train_epochs=1",
         "training.learning_rate=1e-5",
-        "training.training_seq_len=64",
+        # Shorter than the dumped sequences, so the offline collator truncates them.
+        "training.training_seq_len=16",
         "training.save_steps=1",
         # torch.compile is smoke-tested once by test_llama_eagle3[1-False]; skip its warmup here.
         "eagle.eagle_use_torch_compile=false",

From faa41c07a4ebca1e80d6b9f4be2c504b2c59cc57 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:07:06 +0000
Subject: [PATCH 27/34] Run the gated Nemotron v3 dataset smoke test only on
 the nightly

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>
---
 tests/gpu/torch/utils/test_dataset_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/gpu/torch/utils/test_dataset_utils.py b/tests/gpu/torch/utils/test_dataset_utils.py
index bc9451ca03a..f83266762b6 100644
--- a/tests/gpu/torch/utils/test_dataset_utils.py
+++ b/tests/gpu/torch/utils/test_dataset_utils.py
@@ -42,6 +42,9 @@ def qwen3_tokenizer():
     return tokenizer
 
 
+@pytest.mark.release(
+    reason="Streams the 7 gated Nemotron v3 datasets; runs on the nightly schedule."
+)
 def test_get_dataset_dataloader_nemotron_v3_chat_template(qwen3_tokenizer):
     if not get_token():
         pytest.skip(

From bdfa096179acdd191b3b5c76c966781742cb707e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:11:11 +0000
Subject: [PATCH 28/34] Drop DataLoader workers from the torch_onnx calibration
 loaders

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>
---
 examples/torch_onnx/torch_quant_to_onnx.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py
index bd935146e0e..a08338d053e 100644
--- a/examples/torch_onnx/torch_quant_to_onnx.py
+++ b/examples/torch_onnx/torch_quant_to_onnx.py
@@ -27,7 +27,6 @@
 
 import timm
 import torch
-import torch.multiprocessing as mp
 import torch.nn.functional as F
 from datasets import load_dataset
 from download_example_onnx import export_to_onnx
@@ -56,9 +55,6 @@
 """
 
 
-mp.set_start_method("spawn", force=True)  # Needed for data loader with multiple workers
-
-
 _FP8_CONV_OVERRIDE: list = [
     {
         "parent_class": "nn.Conv2d",
@@ -311,16 +307,17 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals
     calib_tensor = [transforms(img) for img in images]
     calib_tensor = [t.to(device) for t in calib_tensor]
 
+    # The samples already sit on ``device``; worker processes would only add spawn and import cost.
     if with_labels:
         labels = dataset["train"][:data_size]["label"]
         labels = torch.tensor(labels, device=device)
         calib_dataset = [{"image": img, "label": lbl} for img, lbl in zip(calib_tensor, labels)]
         return torch.utils.data.DataLoader(
-            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=0
         )
     else:
         return torch.utils.data.DataLoader(
-            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=0
         )
 
 

From 2815b82d547b3300ad42565496bb9b1fc6ebd5ad Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:11:23 +0000
Subject: [PATCH 29/34] Skip swinv2 int8/mxfp8/auto on PRs and run trtexec for
 one iteration

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>
---
 examples/torch_onnx/torch_quant_to_onnx.py            | 4 ++++
 tests/examples/torch_onnx/test_torch_quant_to_onnx.py | 7 +++++++
 2 files changed, 11 insertions(+)

diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py
index a08338d053e..a9a441e1a6d 100644
--- a/examples/torch_onnx/torch_quant_to_onnx.py
+++ b/examples/torch_onnx/torch_quant_to_onnx.py
@@ -758,6 +758,10 @@ def build_trt_engine(onnx_path, optimization_level=4):
         f"--onnx={onnx_path}",
         "--stronglyTyped",
         f"--builderOptimizationLevel={optimization_level}",
+        # One inference iteration shows the engine runs; the default 3 s timed benchmark is not shown.
+        "--warmUp=0",
+        "--duration=0",
+        "--iterations=1",
     ]
     print(f"\nBuilding TensorRT engine: {' '.join(cmd)}")
     try:
diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py
index 28de24e9429..9fa4aa958d1 100644
--- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py
+++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py
@@ -15,6 +15,7 @@
 
 
 import importlib
+import os
 from collections import defaultdict
 from pathlib import Path
 
@@ -38,6 +39,10 @@
     "resnet50": ("resnet50", None),
 }
 
+# PRs run swinv2 for fp8 and nvfp4 only (swin_tiny covers every format); scheduled runs keep all.
+_PR_CI = os.environ.get("GITHUB_REF", "").startswith("refs/heads/pull-request/")
+_SKIPPED_ON_PR = {("swinv2_tiny", "int8"), ("swinv2_tiny", "mxfp8"), ("swinv2_tiny", "auto")}
+
 
 def _assert_residual_inputs_are_quantized(onnx_save_path):
     model = onnx.load(onnx_save_path)
@@ -74,6 +79,8 @@ def _assert_residual_inputs_are_quantized(onnx_save_path):
 def test_torch_onnx(tmp_path, model_key, qformat):
     if model_key == "resnet50" and qformat not in _RESNET_RECIPE_QFORMATS:
         pytest.skip("Only FP8 and INT8 quantization are supported for ResNet")
+    if _PR_CI and (model_key, qformat) in _SKIPPED_ON_PR:
+        pytest.skip("Covered by swin_tiny on PRs; runs in the scheduled CI")
 
     timm_model_name, model_kwargs = _MODELS[model_key]
     onnx_save_path = tmp_path / f"{model_key}.{qformat}.onnx"

From 8a1b4659080b0de22e95eaa552aefebe753d2b84 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:14:01 +0000
Subject: [PATCH 30/34] Add a worker runner so example steps share one python
 launch

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>
---
 tests/_test_utils/examples/example_runner.py | 425 +++++++++++++
 tests/_test_utils/examples/run_command.py    |  33 +-
 tests/unit/test_example_runner.py            | 590 +++++++++++++++++++
 3 files changed, 1039 insertions(+), 9 deletions(-)
 create mode 100644 tests/_test_utils/examples/example_runner.py
 create mode 100644 tests/unit/test_example_runner.py

diff --git a/tests/_test_utils/examples/example_runner.py b/tests/_test_utils/examples/example_runner.py
new file mode 100644
index 00000000000..f303d861d92
--- /dev/null
+++ b/tests/_test_utils/examples/example_runner.py
@@ -0,0 +1,425 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run ``python <script>.py`` example steps in one long-lived worker instead of one launch per step.
+
+A fresh interpreter spends 10-20 s importing torch, transformers and modelopt before an example does any
+work, and an example test runs several steps. The worker pays that once. Each step runs like
+``python script.py`` (``runpy`` as ``__main__``, ``sys.argv``, the example directory as cwd and first on
+``sys.path``); afterwards everything it changed in the process is put back, so the next step starts clean
+while the imports stay cached. Environment variables are applied to ``os.environ``, so those read at
+interpreter start-up or at import time keep the worker's values.
+
+The worker is a separate process, not the pytest process, so a hung step can be killed (that step fails and
+the next one gets a fresh worker) and pytest holds no CUDA context. The worker is also recycled after any
+failed step. ``run_command.run_example_command`` hands an installed ``ExampleRunner`` the steps it
+``accepts``; every other step, and every step when ``MODELOPT_EXAMPLE_RUNNER=subprocess``, still runs as a
+subprocess. ``MODELOPT_EXAMPLE_RUNNER=inprocess`` runs the steps in the calling process instead: no
+isolation, but ``pdb`` works.
+"""
+
+import contextlib
+import gc
+import importlib
+import logging
+import multiprocessing
+import os
+import random
+import runpy
+import shutil
+import signal
+import subprocess
+import sys
+import tempfile
+import traceback
+import warnings
+from collections.abc import Callable, Iterable
+from pathlib import Path
+
+EXAMPLES_DIR = Path(__file__).parents[3] / "examples"
+
+_SIGNALS = [getattr(signal, n) for n in ("SIGTERM", "SIGINT", "SIGHUP") if hasattr(signal, n)]
+# A clean worker exit lets its atexit handlers run, which is how coverage writes its data.
+_WORKER_EXIT_TIMEOUT_S = 30
+_KILL_TIMEOUT_S = 5
+
+
+class StepFailed(subprocess.CalledProcessError):
+    """A step exited non-zero, raised, or took its worker down; ``output`` is what it printed."""
+
+    def __init__(self, returncode: int, cmd: list[str], output: str):
+        super().__init__(returncode, cmd, output=output)
+        # What run_example_command scans for transient Hub errors.
+        self.captured_output = output
+
+
+class StepTimeout(subprocess.TimeoutExpired):
+    """A step outlived its timeout and its worker was killed; ``output`` is what it had printed."""
+
+    def __init__(self, cmd: list[str], timeout: float, output: str):
+        super().__init__(cmd, timeout, output=output)
+        self.captured_output = output
+
+
+def _loggers() -> list[logging.Logger]:
+    with logging._lock:  # a getLogger() on another thread mutates loggerDict
+        named = [
+            lg for lg in logging.root.manager.loggerDict.values() if isinstance(lg, logging.Logger)
+        ]
+    return [logging.root, *named]
+
+
+def _torch_flags(torch) -> list[tuple]:
+    """The (object, attribute) pairs a script may flip, e.g. onnx_ptq/evaluation.py sets cudnn's."""
+    backends = torch.backends
+    return [
+        (backends.cudnn, "deterministic"),
+        (backends.cudnn, "benchmark"),
+        (backends.cudnn, "allow_tf32"),
+        (backends.cuda.matmul, "allow_tf32"),
+    ]
+
+
+def _purge_example_modules(before: set[str], examples_dir: Path) -> set[str]:
+    """Drop modules the step imported from ``examples/``: ``utils``, ``config``, sibling example dirs.
+
+    The next step may bring a different module of the same name. Third-party and modelopt modules
+    stay cached, which is the saving.
+    """
+    root = examples_dir.resolve()
+    purged = set()
+    for name in set(sys.modules) - before:
+        file = getattr(sys.modules[name], "__file__", None)
+        if file and Path(file).resolve().is_relative_to(root):
+            purged.add(name)
+            del sys.modules[name]
+    return purged
+
+
+def _restore_loggers(snapshot: dict, purged: set[str]) -> None:
+    """Restore levels everywhere, handler lists only where the script owns them.
+
+    That is root, ``__main__`` and loggers named after purged modules: a library may configure its
+    own logger lazily during a step and must keep its handler.
+    """
+    owned = {"__main__", *(name.split(".")[0] for name in purged)}
+    for lg in _loggers():
+        is_owned = lg is logging.root or lg.name.split(".")[0] in owned
+        if lg in snapshot:
+            handlers, level, propagate, disabled = snapshot[lg]
+            lg.setLevel(level)
+            lg.propagate, lg.disabled = propagate, disabled
+            if is_owned:
+                lg.handlers[:] = handlers
+        elif is_owned:
+            lg.handlers[:] = []
+            lg.setLevel(logging.NOTSET)
+
+
+@contextlib.contextmanager
+def _restored_state(examples_dir: Path):
+    """Put back what a step changed in the process, which ``python script.py`` discards on exit."""
+    argv, path, cwd, environ = list(sys.argv), list(sys.path), os.getcwd(), dict(os.environ)
+    modules = set(sys.modules)
+    signal_handlers = {sig: signal.getsignal(sig) for sig in _SIGNALS}
+    loggers = {lg: (list(lg.handlers), lg.level, lg.propagate, lg.disabled) for lg in _loggers()}
+    log_disable = logging.root.manager.disable
+    start_method = multiprocessing.get_start_method(allow_none=True)
+    py_random = random.getstate()
+    numpy = sys.modules.get("numpy")
+    np_random = numpy.random.get_state() if numpy else None
+    torch = sys.modules.get("torch")
+    torch_rng = torch.get_rng_state() if torch else None
+    torch_flags = (
+        [(obj, name, getattr(obj, name)) for obj, name in _torch_flags(torch)] if torch else []
+    )
+    with warnings.catch_warnings():
+        try:
+            yield
+        finally:
+            sys.argv[:] = argv
+            sys.path[:] = path
+            os.chdir(cwd)
+            os.environ.clear()
+            os.environ.update(environ)
+            for sig, handler in signal_handlers.items():
+                # A handler installed from C reads back as None and cannot be restored.
+                if handler is not None:
+                    signal.signal(sig, handler)
+            logging.disable(log_disable)
+            if multiprocessing.get_start_method(allow_none=True) != start_method:
+                multiprocessing.set_start_method(start_method, force=True)
+            random.setstate(py_random)
+            if np_random is not None:
+                numpy.random.set_state(np_random)
+            if torch is not None:
+                torch.set_rng_state(torch_rng)
+                for obj, name, value in torch_flags:
+                    setattr(obj, name, value)
+            _restore_loggers(loggers, _purge_example_modules(modules, examples_dir))
+            gc.collect()
+            if torch is not None and torch.cuda.is_initialized():
+                torch.cuda.empty_cache()
+
+
+def _drain(sink) -> str:
+    sink.flush()
+    sink.seek(0)
+    return sink.read()
+
+
+def run_step(
+    cmd: list[str],
+    cwd: Path | str,
+    env: dict[str, str] | None,
+    sink,
+    hooks: Iterable[Callable] = (),
+    examples_dir: Path = EXAMPLES_DIR,
+) -> tuple[int, str]:
+    """Run ``cmd`` (``[python, script, *args]``) in this process as ``python script ...`` in ``cwd``.
+
+    Everything written to fds 1 and 2, by this process or its children, goes to ``sink``. ``env``
+    replaces ``os.environ`` for the step. Each of ``hooks`` is a callable returning a context
+    manager that wraps the step, for global state the generic restore does not know about.
+    Returns ``(returncode, output)``. Only a ``BaseException`` that is not a script failure, such
+    as ``KeyboardInterrupt`` or a pytest timeout, propagates, with ``captured_output`` set.
+    """
+    script = Path(cwd, cmd[1])
+    sink.seek(0)
+    sink.truncate()
+    sys.stdout.flush()
+    sys.stderr.flush()
+    saved_fds = (os.dup(1), os.dup(2))
+    saved_streams = (sys.stdout, sys.stderr)
+    returncode = 0
+    try:
+        with _restored_state(examples_dir), contextlib.ExitStack() as stack:
+            try:
+                # Child processes and C extensions write to the fds; pytest's capture replaces
+                # sys.stdout and sys.stderr with objects over its own files, so swap those too.
+                os.dup2(sink.fileno(), 1)
+                os.dup2(sink.fileno(), 2)
+                sys.stdout = sys.stderr = sink
+                os.chdir(cwd)
+                if env is not None:
+                    os.environ.clear()
+                    os.environ.update(env)
+                sys.argv[:] = [str(script), *cmd[2:]]
+                sys.path.insert(0, str(script.parent))
+                importlib.invalidate_caches()
+                for hook in hooks:
+                    stack.enter_context(hook())
+                runpy.run_path(str(script), run_name="__main__")
+            except SystemExit as e:  # argparse errors, ``sys.exit(1)`` after a logged failure
+                if e.code not in (0, None):
+                    returncode = e.code if isinstance(e.code, int) else 1
+                    if not isinstance(e.code, int):
+                        print(e.code, file=sys.stderr)
+            except Exception:
+                traceback.print_exc()
+                returncode = 1
+    except BaseException as e:
+        e.captured_output = _drain(sink)
+        raise
+    finally:
+        sys.stdout, sys.stderr = saved_streams
+        os.dup2(saved_fds[0], 1)
+        os.dup2(saved_fds[1], 2)
+        os.close(saved_fds[0])
+        os.close(saved_fds[1])
+    return returncode, _drain(sink)
+
+
+def _preload(modules: Iterable[str]) -> None:
+    """Import what ``_restored_state`` snapshots, so a first step cannot change it unnoticed."""
+    for name in modules:
+        with contextlib.suppress(ImportError):
+            importlib.import_module(name)
+
+
+def _open_sink(path: str):
+    return open(path, "w+", encoding="utf-8", errors="replace", buffering=1)
+
+
+def _worker_main(conn, sink_path, hooks, examples_dir, preload) -> None:
+    # Own session and process group: killing the worker takes a step's children along.
+    os.setsid()
+    # Steps may start DataLoader workers or Pools, which a daemonic process is not allowed to.
+    multiprocessing.current_process().daemon = False
+    _preload(preload)
+    sink = _open_sink(sink_path)
+    while True:
+        try:
+            request = conn.recv()
+        except EOFError:
+            return
+        if request is None:
+            return
+        cmd, cwd, env = request
+        try:
+            reply = run_step(cmd, cwd, env, sink, hooks, Path(examples_dir))
+        except BaseException:
+            reply = (1, traceback.format_exc())
+        conn.send(reply)
+
+
+class ExampleRunner:
+    """Run the example steps listed in ``scripts`` in a long-lived worker process.
+
+    Install it with ``run_command.set_in_process_runner``. ``scripts`` are the names of the scripts
+    (as in ``["python", name, ...]``) known to be safe to run repeatedly in one process. ``hooks``
+    wrap every step, see ``run_step``; the worker imports them, so they must be module-level.
+    ``timeout`` bounds one step; by default only pytest-timeout does, and it kills the worker too.
+    ``preload`` names modules imported before the first step.
+    """
+
+    def __init__(
+        self,
+        scripts: Iterable[str],
+        hooks: Iterable[Callable] = (),
+        examples_dir: Path | str = EXAMPLES_DIR,
+        timeout: float | None = None,
+        preload: Iterable[str] = ("torch",),
+        in_process: bool | None = None,
+    ):
+        self.scripts = frozenset(scripts)
+        self._hooks = tuple(hooks)
+        self._examples_dir = Path(examples_dir)
+        self._timeout = timeout
+        self._preload = tuple(preload)
+        if in_process is None:
+            in_process = os.environ.get("MODELOPT_EXAMPLE_RUNNER") == "inprocess"
+        self._in_process = in_process
+        self._sink_dir = tempfile.mkdtemp(prefix="example_runner_")
+        self._sink_path = os.path.join(self._sink_dir, "output.txt")
+        # In-process only, and never closed: a handler built during a step may still hold it.
+        self._sink = None
+        self._proc = None
+        self._conn = None
+
+    def accepts(self, cmd_parts: list[str]) -> bool:
+        """Whether the step is ``python <one of scripts> ...``."""
+        interpreters = {"python", "python3", Path(sys.executable).name}
+        return (
+            len(cmd_parts) > 1
+            and Path(str(cmd_parts[0])).name in interpreters
+            and str(cmd_parts[1]) in self.scripts
+        )
+
+    def __call__(
+        self,
+        cmd_parts: list[str],
+        example_path: str,
+        env: dict[str, str] | None = None,
+        timeout: float | None = None,
+    ) -> str:
+        """Run a step in ``examples/<example_path>`` and return its output.
+
+        Raises ``StepFailed`` or ``StepTimeout``. ``env`` is the step's whole environment; ``None``
+        keeps the worker's own.
+        """
+        cmd = [str(part) for part in cmd_parts]
+        cwd = self._examples_dir / example_path
+        timeout = self._timeout if timeout is None else timeout
+        output = ""
+        try:
+            if self._in_process:
+                output = self._run_in_process(cmd, cwd, env)
+            else:
+                output = self._run_in_worker(cmd, cwd, env, timeout)
+            return output
+        except BaseException as e:
+            output = getattr(e, "captured_output", "")
+            raise
+        finally:
+            print(output)  # keep the step's output in the test log
+
+    def close(self) -> None:
+        self._stop_worker()
+        shutil.rmtree(self._sink_dir, ignore_errors=True)
+
+    def _run_in_process(self, cmd: list[str], cwd: Path, env: dict[str, str] | None) -> str:
+        if self._sink is None:
+            _preload(self._preload)
+            self._sink = _open_sink(self._sink_path)
+        returncode, output = run_step(cmd, cwd, env, self._sink, self._hooks, self._examples_dir)
+        if returncode != 0:
+            raise StepFailed(returncode, cmd, output)
+        return output
+
+    def _run_in_worker(
+        self, cmd: list[str], cwd: Path, env: dict[str, str] | None, timeout: float | None
+    ) -> str:
+        self._start_worker()
+        try:
+            self._conn.send((cmd, str(cwd), env))
+            reply = self._conn.recv() if self._conn.poll(timeout) else None
+        except (EOFError, OSError):
+            # The worker died during the step: segfault, OOM kill.
+            output = self._partial_output()
+            raise StepFailed(self._kill_worker(), cmd, output) from None
+        except BaseException as e:
+            # Interrupted while waiting, e.g. by pytest-timeout: do not leave the step running.
+            e.captured_output = self._partial_output()
+            self._kill_worker()
+            raise
+        if reply is None:
+            # Hung: only this step fails, the next one gets a fresh worker.
+            output = self._partial_output()
+            self._kill_worker()
+            raise StepTimeout(cmd, timeout, output)
+        returncode, output = reply
+        if returncode != 0:
+            self._stop_worker()  # do not trust the state a failed step left behind
+            raise StepFailed(returncode, cmd, output)
+        return output
+
+    def _start_worker(self) -> None:
+        if self._proc is not None and self._proc.is_alive():
+            return
+        context = multiprocessing.get_context("spawn")
+        self._conn, child_conn = context.Pipe()
+        args = (child_conn, self._sink_path, self._hooks, str(self._examples_dir), self._preload)
+        self._proc = context.Process(target=_worker_main, args=args, daemon=True)
+        self._proc.start()
+        child_conn.close()  # only the worker holds this end, so its death shows up as EOF
+
+    def _partial_output(self) -> str:
+        try:
+            return Path(self._sink_path).read_text(encoding="utf-8", errors="replace")
+        except OSError:
+            return ""
+
+    def _kill_worker(self) -> int:
+        """SIGKILL the worker's process group and return the worker's exit code."""
+        proc, self._proc = self._proc, None
+        with contextlib.suppress(ProcessLookupError, PermissionError):
+            # Before the pid is reaped, so it cannot have been recycled.
+            os.killpg(proc.pid, signal.SIGKILL)
+        proc.join(_KILL_TIMEOUT_S)
+        return -1 if proc.exitcode is None else proc.exitcode
+
+    def _stop_worker(self) -> None:
+        """Let the worker exit on its own, so its coverage data is written, or kill it."""
+        proc = self._proc
+        if proc is None:
+            return
+        if proc.is_alive():
+            with contextlib.suppress(OSError):
+                self._conn.send(None)
+            proc.join(_WORKER_EXIT_TIMEOUT_S)
+        if proc.is_alive():
+            self._kill_worker()
+        else:
+            self._proc = None
diff --git a/tests/_test_utils/examples/run_command.py b/tests/_test_utils/examples/run_command.py
index 218106f9f89..f7aee9d95f8 100644
--- a/tests/_test_utils/examples/run_command.py
+++ b/tests/_test_utils/examples/run_command.py
@@ -137,13 +137,20 @@ def _drain():
     return returncode, "".join(chunks)
 
 
-# Set by a suite's conftest to run a command without spawning a subprocess (see
-# ``_test_utils.examples.megatron_example_runner``). Unset means always use a subprocess.
+# Set by a suite's conftest to run steps without a subprocess each. Unset means always use one, and so
+# does MODELOPT_EXAMPLE_RUNNER=subprocess (to bisect a suspected leak between steps).
 _in_process_runner = None
 
 
 def set_in_process_runner(runner) -> None:
-    """Register ``runner(cmd_parts, example_path) -> str``; it must not fall back to a subprocess."""
+    """Register either of two kinds of runner.
+
+    A plain callable ``runner(cmd_parts, example_path) -> str`` runs every step in the ambient
+    environment and must not fall back to a subprocess (``megatron_example_runner``). An object
+    with ``accepts(cmd_parts)`` and ``runner(cmd_parts, example_path, env=env) -> str`` takes the
+    step's environment and only the steps it accepts; the rest run as subprocesses
+    (``example_runner.ExampleRunner``).
+    """
     global _in_process_runner
     _in_process_runner = runner
 
@@ -158,24 +165,32 @@ def run_example_command(
 ) -> str | None:
     """Run an example command, retrying transient HuggingFace access errors."""
     print(f"[{example_path}] Running command: {cmd_parts}")
-    in_process = _in_process_runner if env is None else None
-    if _in_process_runner is not None and env is not None:
-        # The in-process runner uses the ambient environment, so a caller-supplied env would be
+    in_process = (
+        None if os.environ.get("MODELOPT_EXAMPLE_RUNNER") == "subprocess" else _in_process_runner
+    )
+    takes_env = hasattr(in_process, "accepts")
+    if takes_env and not in_process.accepts(cmd_parts):
+        in_process = None
+    elif in_process is not None and not takes_env and env is not None:
+        # A plain in-process runner uses the ambient environment, so a caller-supplied env would be
         # silently dropped. Fall back to a subprocess rather than run with the wrong environment.
         warnings.warn(f"[{example_path}] env= given; running this step as a subprocess")
+        in_process = None
     env = os.environ.copy() if env is None else env
     cwd = MODELOPT_ROOT / "examples" / example_path
 
     for attempt in range(hf_max_retries + 1):
         if setup_free_port:
-            # Subprocess steps only: an in-process runner picks its own free port, since env is a
-            # copy it never sees.
+            # A plain in-process runner picks its own free port, since env is a copy it never sees.
             env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
         if in_process is not None:
             # Inside the loop so in-process steps get the same transient-HuggingFace retries;
             # these tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
             try:
-                result = in_process(cmd_parts, example_path)
+                if takes_env:
+                    result = in_process(cmd_parts, example_path, env=env)
+                else:
+                    result = in_process(cmd_parts, example_path)
             except Exception as e:
                 # Re-raise unless it looks transient, so a real failure keeps its traceback
                 # instead of being flattened into CalledProcessError.
diff --git a/tests/unit/test_example_runner.py b/tests/unit/test_example_runner.py
new file mode 100644
index 00000000000..65d03dad3b4
--- /dev/null
+++ b/tests/unit/test_example_runner.py
@@ -0,0 +1,590 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for ``ExampleRunner``: example steps run in a long-lived worker and leave no state behind."""
+
+import contextlib
+import json
+import os
+import re
+import signal
+import subprocess
+import sys
+import textwrap
+import threading
+import time
+from pathlib import Path
+
+import pytest
+from _test_utils.examples import run_command
+from _test_utils.examples.example_runner import ExampleRunner, StepFailed, StepTimeout
+
+pytestmark = pytest.mark.usefixtures("skip_on_windows")
+
+
+class _Lane:
+    """One example directory of a runner, with helpers to write and run its scripts."""
+
+    def __init__(self, runner, directory):
+        directory.mkdir(parents=True)
+        self.runner, self.dir = runner, directory
+
+    def write(self, name, body):
+        (self.dir / name).write_text(textwrap.dedent(body))
+
+    def run(self, script, *args, **kwargs):
+        return self.runner(["python", script, *args], self.dir.name, **kwargs)
+
+    def probe(self, script="probe.py"):
+        return json.loads(self.run(script).strip().splitlines()[-1])
+
+
+def _is_alive(pid):
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    try:  # a zombie is dead; nothing reaps it where pid 1 does not
+        return Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split()[0] != "Z"
+    except OSError:
+        return True
+
+
+def _wait_until_dead(pid):
+    for _ in range(100):
+        if not _is_alive(pid):
+            return
+        time.sleep(0.1)
+    os.kill(pid, signal.SIGKILL)
+    pytest.fail(f"process {pid} outlived its step")
+
+
+@contextlib.contextmanager
+def _in_process_lane(tmp_path, hooks=()):
+    runner = ExampleRunner((), hooks, tmp_path, preload=(), in_process=True)
+    try:
+        yield _Lane(runner, tmp_path / "example")
+    finally:
+        runner.close()
+
+
+def _lane_factory(runner, examples_dir, request):
+    name = re.sub(r"\W+", "_", request.node.name)
+    return lambda suffix="": _Lane(runner, examples_dir / f"{name}{suffix}")
+
+
+@pytest.fixture(scope="module")
+def examples_dir(tmp_path_factory):
+    return tmp_path_factory.mktemp("examples")
+
+
+@pytest.fixture(scope="module")
+def worker_runner(examples_dir):
+    # No preload: every failed step recycles the worker, and importing torch takes seconds.
+    runner = ExampleRunner((), examples_dir=examples_dir, preload=(), in_process=False)
+    yield runner
+    runner.close()
+
+
+@pytest.fixture(scope="module", params=["worker", "in_process"])
+def backend(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def any_runner(backend, examples_dir, worker_runner):
+    if backend == "worker":
+        yield worker_runner
+        return
+    runner = ExampleRunner((), examples_dir=examples_dir, preload=(), in_process=True)
+    yield runner
+    runner.close()
+
+
+@pytest.fixture
+def make_lane(worker_runner, examples_dir, request):
+    """Lanes of the shared worker."""
+    return _lane_factory(worker_runner, examples_dir, request)
+
+
+@pytest.fixture
+def make_lane_on_both_backends(any_runner, examples_dir, request):
+    """Lanes of the shared worker, and of an in-process runner, as the test runs once per backend."""
+    return _lane_factory(any_runner, examples_dir, request)
+
+
+@pytest.fixture
+def worker_lane(tmp_path):
+    """A lane of a worker of its own, for tests that kill or stop it."""
+    runner = ExampleRunner((), examples_dir=tmp_path, preload=(), in_process=False)
+    yield _Lane(runner, tmp_path / "example")
+    runner.close()
+
+
+def test_step_runs_like_python_script_py(make_lane_on_both_backends):
+    lane = make_lane_on_both_backends()
+    lane.write(
+        "s.py",
+        """
+        import argparse, os, sys
+
+        def main():
+            parser = argparse.ArgumentParser()
+            parser.add_argument("--x", type=int)
+            cwd, argv0 = os.path.basename(os.getcwd()), os.path.basename(sys.argv[0])
+            print("x", parser.parse_args().x, "cwd", cwd, "argv0", argv0)
+
+        if __name__ == "__main__":
+            main()
+        """,
+    )
+    assert f"x 3 cwd {lane.dir.name} argv0 s.py" in lane.run("s.py", "--x", "3")
+
+
+def test_failed_step_raises_step_failed_and_the_runner_recovers(make_lane_on_both_backends):
+    lane = make_lane_on_both_backends()
+    lane.write("bad.py", "raise ValueError('bad')")
+    lane.write("ok.py", "print('still fine')")
+    with pytest.raises(StepFailed) as failure:
+        lane.run("bad.py")
+    assert isinstance(failure.value, subprocess.CalledProcessError)
+    assert failure.value.returncode == 1
+    assert "ValueError: bad" in failure.value.output
+    assert failure.value.captured_output == failure.value.output
+    assert "still fine" in lane.run("ok.py")
+
+
+@pytest.mark.parametrize(
+    ("body", "returncode", "text"),
+    [
+        ("import sys; sys.exit(3)", 3, ""),
+        ("import sys; sys.exit('boom')", 1, "boom"),
+        ("import argparse; argparse.ArgumentParser().parse_args(['--nope'])", 2, "unrecognized"),
+    ],
+)
+def test_exit_statuses_map_to_the_return_code(make_lane, body, returncode, text):
+    lane = make_lane()
+    lane.write("bad.py", body)
+    with pytest.raises(StepFailed) as failure:
+        lane.run("bad.py")
+    assert failure.value.returncode == returncode
+    assert text in failure.value.output
+
+
+def test_sys_exit_zero_is_success(make_lane):
+    lane = make_lane()
+    lane.write("z.py", "import sys; print('done'); sys.exit(0)")
+    assert "done" in lane.run("z.py")
+
+
+def test_process_state_changed_by_a_step_is_restored(make_lane_on_both_backends):
+    lane = make_lane_on_both_backends()
+    lane.write(
+        "leak.py",
+        """
+        import logging, os, signal, sys, warnings
+
+        os.environ["LEAK"] = "1"
+        os.chdir("/")
+        sys.path.insert(0, "/nonexistent-leak")
+        signal.signal(signal.SIGTERM, lambda *args: None)
+        logger = logging.getLogger("__main__")
+        logger.addHandler(logging.StreamHandler(sys.stdout))
+        logger.setLevel(10)
+        logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+        logging.disable(logging.CRITICAL)
+        warnings.simplefilter("ignore")
+        sys.argv.append("leak")
+        """,
+    )
+    lane.write(
+        "probe.py",
+        """
+        import json, logging, os, signal, sys, warnings
+
+        print(json.dumps({
+            "leak": "LEAK" in os.environ,
+            "path": "/nonexistent-leak" in sys.path,
+            "argv": len(sys.argv),
+            "sigterm": str(signal.getsignal(signal.SIGTERM)),
+            "main_handlers": len(logging.getLogger("__main__").handlers),
+            "root_handlers": len(logging.getLogger().handlers),
+            "root_level": logging.getLogger().level,
+            "disabled": logging.root.manager.disable,
+            "filters": len(warnings.filters),
+        }))
+        """,
+    )
+    cwd, before = os.getcwd(), lane.probe()
+    lane.run("leak.py")
+    after = lane.probe()
+    assert after == before
+    assert after["leak"] is False
+    assert after["main_handlers"] == 0
+    assert os.getcwd() == cwd
+
+
+def test_random_state_start_method_and_backend_flags_are_restored(backend, examples_dir):
+    # Preload torch up front, as the real lanes do.
+    runner = ExampleRunner(
+        (), examples_dir=examples_dir, preload=("torch",), in_process=backend == "in_process"
+    )
+    lane = _Lane(runner, examples_dir / f"state_{backend}")
+    lane.write(
+        "mutate.py",
+        """
+        import multiprocessing, random
+        import numpy, torch
+
+        multiprocessing.set_start_method("forkserver", force=True)
+        random.seed(1)
+        numpy.random.seed(1)
+        torch.manual_seed(1)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = True
+        """,
+    )
+    lane.write(
+        "probe.py",
+        """
+        import json, multiprocessing, random
+        import numpy, torch
+
+        print(json.dumps({
+            "start_method": multiprocessing.get_start_method(allow_none=True),
+            "python": random.random(),
+            "numpy": float(numpy.random.rand()),
+            "torch": float(torch.rand(1)),
+            "deterministic": torch.backends.cudnn.deterministic,
+            "benchmark": torch.backends.cudnn.benchmark,
+        }))
+        """,
+    )
+    try:
+        before = lane.probe()
+        lane.run("mutate.py")
+        assert lane.probe() == before
+    finally:
+        runner.close()
+
+
+def test_handlers_added_by_each_step_do_not_accumulate(make_lane):
+    # What diffusers' quantize.py does with its setup_logging().
+    lane = make_lane()
+    lane.write(
+        "logs.py",
+        """
+        import logging, sys
+
+        log = logging.getLogger(__name__)
+        log.setLevel(logging.INFO)
+        log.addHandler(logging.StreamHandler(sys.stdout))
+        log.info("LINE")
+        """,
+    )
+    for _ in range(3):
+        assert lane.run("logs.py").count("LINE") == 1
+
+
+def test_modules_of_the_example_dir_are_not_shared_but_imports_stay_cached(make_lane):
+    first, second = make_lane("_a"), make_lane("_b")
+    for lane, value in ((first, "A"), (second, "B")):
+        lane.write("utils.py", f"VALUE = '{value}'")
+        lane.write(
+            "main.py", "import sys, utils\nprint('got', utils.VALUE, 'colorsys' in sys.modules)"
+        )
+    first.write("warm.py", "import colorsys")
+    first.run("warm.py")
+    seen = [lane.run("main.py") for lane in (first, second, first)]
+    assert ["got A True" in seen[0], "got B True" in seen[1], "got A True" in seen[2]] == [True] * 3
+
+
+def test_env_is_the_whole_environment_of_the_step_and_is_restored(make_lane):
+    lane = make_lane()
+    lane.write("keys.py", "import os; print('KEYS', sorted(os.environ))")
+    assert "KEYS ['ONLY']" in lane.run("keys.py", env={"ONLY": "1"})
+    assert "ONLY" not in lane.run("keys.py")
+
+
+def test_output_of_fds_and_child_processes_is_captured(make_lane):
+    lane = make_lane()
+    lane.write(
+        "o.py",
+        """
+        import os, subprocess, sys
+
+        print("py")
+        print("err", file=sys.stderr)
+        os.write(1, b"fd1\\n")
+        subprocess.run(["echo", "child"])
+        """,
+    )
+    output = lane.run("o.py")
+    assert all(text in output for text in ("py", "err", "fd1", "child"))
+
+
+def test_step_can_start_child_processes(make_lane):
+    # DataLoader workers and Pools: a daemonic worker would refuse.
+    lane = make_lane()
+    lane.write(
+        "fork.py",
+        """
+        import multiprocessing
+
+        def child():
+            print("child ran")
+
+        if __name__ == "__main__":
+            process = multiprocessing.get_context("fork").Process(target=child)
+            process.start()
+            process.join()
+            print("exit", process.exitcode)
+        """,
+    )
+    output = lane.run("fork.py")
+    assert "child ran" in output
+    assert "exit 0" in output
+
+
+def test_hooks_wrap_every_step(tmp_path):
+    events = []
+
+    @contextlib.contextmanager
+    def hook():
+        events.append("enter")
+        try:
+            yield
+        finally:
+            events.append("exit")
+
+    with _in_process_lane(tmp_path, (hook,)) as lane:
+        lane.write("s.py", "print('ran')")
+        assert "ran" in lane.run("s.py")
+    assert events == ["enter", "exit"]
+
+
+def test_failing_hook_fails_the_step(tmp_path):
+    def hook():
+        raise RuntimeError("hook failed")
+
+    with _in_process_lane(tmp_path, (hook,)) as lane:
+        lane.write("s.py", "print('ran')")
+        with pytest.raises(StepFailed) as failure:
+            lane.run("s.py")
+    assert "hook failed" in failure.value.output
+    assert "ran" not in failure.value.output
+
+
+def test_worker_is_reused_until_a_step_fails(worker_lane):
+    worker_lane.write("pid.py", "import os; print('pid', os.getpid())")
+    worker_lane.write("bad.py", "raise SystemExit(1)")
+
+    def pid():
+        return worker_lane.run("pid.py").split()[-1]
+
+    first = pid()
+    assert pid() == first != str(os.getpid())
+    with pytest.raises(StepFailed):
+        worker_lane.run("bad.py")
+    assert pid() != first
+
+
+def test_timeout_kills_the_hung_step_with_its_children_and_the_next_step_gets_a_fresh_worker(
+    worker_lane,
+):
+    worker_lane.write("pid.py", "import os; print('pid', os.getpid())")
+    worker_lane.write(
+        "hang.py",
+        """
+        import subprocess, time
+
+        child = subprocess.Popen(["sleep", "60"])
+        print("child", child.pid, flush=True)
+        time.sleep(60)
+        """,
+    )
+    first = worker_lane.run("pid.py")
+    started = time.monotonic()
+    with pytest.raises(StepTimeout) as timeout:
+        worker_lane.run("hang.py", timeout=2)
+    assert time.monotonic() - started < 15
+    child = int(re.search(r"child (\d+)", timeout.value.captured_output).group(1))
+    _wait_until_dead(child)
+    assert worker_lane.run("pid.py") != first
+
+
+@pytest.mark.parametrize(
+    ("death", "returncode"),
+    [("os._exit(7)", 7), ("os.kill(os.getpid(), signal.SIGKILL)", -signal.SIGKILL)],
+)
+def test_worker_death_fails_the_step_not_the_session(worker_lane, death, returncode):
+    worker_lane.write("die.py", f"import os, signal\nprint('about to die', flush=True)\n{death}")
+    worker_lane.write("ok.py", "print('alive')")
+    with pytest.raises(StepFailed) as failure:
+        worker_lane.run("die.py")
+    assert failure.value.returncode == returncode
+    assert "about to die" in failure.value.output
+    assert "alive" in worker_lane.run("ok.py")
+
+
+def test_interrupt_while_waiting_kills_the_step(worker_lane):
+    # What pytest-timeout's signal method does: raise in the main thread while the runner waits.
+    class Interrupted(BaseException):
+        pass
+
+    def interrupt(*_):
+        raise Interrupted
+
+    worker_lane.write(
+        "hang.py",
+        """
+        import os, pathlib, time
+
+        pathlib.Path(__file__).with_name("pid").write_text(str(os.getpid()))
+        print("before hang", flush=True)
+        time.sleep(60)
+        """,
+    )
+    previous = signal.signal(signal.SIGUSR1, interrupt)
+    main_thread = threading.main_thread().ident
+    timer = threading.Timer(1.0, signal.pthread_kill, (main_thread, signal.SIGUSR1))
+    timer.start()
+    try:
+        with pytest.raises(Interrupted) as interrupted:
+            worker_lane.run("hang.py", timeout=30)
+    finally:
+        timer.cancel()
+        signal.signal(signal.SIGUSR1, previous)
+    assert "before hang" in interrupted.value.captured_output
+    _wait_until_dead(int((worker_lane.dir / "pid").read_text()))
+
+
+def test_close_lets_the_worker_exit_cleanly(worker_lane):
+    # A clean exit is what lets coverage write the worker's data.
+    worker_lane.write(
+        "marker.py",
+        """
+        import atexit, pathlib
+
+        marker = pathlib.Path(__file__).with_name("exited")
+        atexit.register(marker.write_text, "clean")
+        """,
+    )
+    worker_lane.run("marker.py")
+    marker = worker_lane.dir / "exited"
+    assert not marker.exists()
+    worker_lane.runner.close()
+    assert marker.read_text() == "clean"
+
+
+def test_accepts_only_listed_scripts_run_with_python():
+    runner = ExampleRunner({"a.py"}, preload=())
+    try:
+        assert runner.accepts(["python", "a.py", "--x", "1"])
+        assert runner.accepts(["python3", "a.py"])
+        assert runner.accepts([sys.executable, "a.py"])
+        assert not runner.accepts(["python", "b.py"])
+        assert not runner.accepts(["bash", "a.py"])
+        assert not runner.accepts(["python", "-m", "a"])
+        assert not runner.accepts(["python"])
+    finally:
+        runner.close()
+
+
+@pytest.fixture
+def installed_runner(tmp_path, monkeypatch):
+    """A worker for ``in_worker.py`` installed in ``run_command``; yields ``examples/demo``."""
+    monkeypatch.setattr(run_command, "MODELOPT_ROOT", tmp_path)
+    demo = tmp_path / "examples" / "demo"
+    demo.mkdir(parents=True)
+    runner = ExampleRunner({"in_worker.py"}, examples_dir=tmp_path / "examples", preload=())
+    run_command.set_in_process_runner(runner)
+    yield demo
+    run_command.set_in_process_runner(None)
+    runner.close()
+
+
+def _run_demo(script, **kwargs):
+    return run_command.run_example_command([sys.executable, script], "demo", **kwargs)
+
+
+def _demo_pid(script):
+    return int(_run_demo(script).split()[-1])
+
+
+def test_run_example_command_gives_listed_scripts_to_the_runner(installed_runner):
+    for name in ("in_worker.py", "elsewhere.py"):
+        (installed_runner / name).write_text("import os; print('pid', os.getpid())")
+    worker = _demo_pid("in_worker.py")
+    assert worker != os.getpid()
+    assert _demo_pid("in_worker.py") == worker
+    assert _demo_pid("elsewhere.py") not in (worker, os.getpid())
+
+
+def test_kill_switch_runs_listed_scripts_as_subprocesses(installed_runner, monkeypatch):
+    (installed_runner / "in_worker.py").write_text("import os; print('pid', os.getpid())")
+    monkeypatch.setenv("MODELOPT_EXAMPLE_RUNNER", "subprocess")
+    assert _demo_pid("in_worker.py") != _demo_pid("in_worker.py")
+
+
+def test_run_example_command_passes_env_to_the_runner(installed_runner):
+    (installed_runner / "in_worker.py").write_text("import os; print('KEYS', sorted(os.environ))")
+    assert "KEYS ['ONLY']" in _run_demo("in_worker.py", env={"ONLY": "1"})
+
+
+def test_failed_step_raises_called_process_error_through_run_example_command(installed_runner):
+    (installed_runner / "in_worker.py").write_text("raise ValueError('nope')")
+    with pytest.raises(subprocess.CalledProcessError) as failure:
+        _run_demo("in_worker.py")
+    assert "ValueError: nope" in failure.value.output
+
+
+def test_transient_hub_errors_are_retried_through_the_runner(installed_runner):
+    (installed_runner / "in_worker.py").write_text(
+        textwrap.dedent(
+            """
+            import pathlib
+
+            marker = pathlib.Path(__file__).with_name("tried")
+            if not marker.exists():
+                marker.write_text("1")
+                raise RuntimeError("ConnectionError: Max retries exceeded")
+            print("recovered")
+            """
+        )
+    )
+    with pytest.warns(UserWarning, match="transient HuggingFace access error"):
+        assert "recovered" in _run_demo("in_worker.py", hf_retry_delay_s=0)
+
+
+def test_plain_runners_keep_their_contract(tmp_path, monkeypatch):
+    monkeypatch.setattr(run_command, "MODELOPT_ROOT", tmp_path)
+    (tmp_path / "examples" / "demo").mkdir(parents=True)
+    calls = []
+
+    def plain_runner(cmd_parts, example_path):
+        calls.append((cmd_parts, example_path))
+        return "ran in process"
+
+    run_command.set_in_process_runner(plain_runner)
+    try:
+        assert run_command.run_example_command(["x.py"], "demo") == "ran in process"
+        with pytest.warns(UserWarning, match="env= given"):
+            output = run_command.run_example_command(
+                [sys.executable, "-c", "print('subprocess')"], "demo", env=os.environ.copy()
+            )
+    finally:
+        run_command.set_in_process_runner(None)
+    assert "subprocess" in output
+    assert calls == [(["x.py"], "demo")]

From d8b9cf3dd26840c3fdf495d870a6ed11087ca5d1 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 00:21:59 +0000
Subject: [PATCH 31/34] Run the torch_onnx lane's torch_quant_to_onnx.py steps
 in one worker

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>
---
 tests/_test_utils/examples/example_runner.py | 17 ++++++++
 tests/examples/torch_onnx/conftest.py        | 44 ++++++++++++++++++++
 tests/unit/test_example_runner.py            | 24 ++++++++++-
 3 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 tests/examples/torch_onnx/conftest.py

diff --git a/tests/_test_utils/examples/example_runner.py b/tests/_test_utils/examples/example_runner.py
index f303d861d92..5895d960b88 100644
--- a/tests/_test_utils/examples/example_runner.py
+++ b/tests/_test_utils/examples/example_runner.py
@@ -241,6 +241,23 @@ def run_step(
     return returncode, _drain(sink)
 
 
+@contextlib.contextmanager
+def keep_post_conversion_plugins():
+    """Hook undoing a step's ``CUSTOM_POST_CONVERSION_PLUGINS.add`` calls.
+
+    torch_quant_to_onnx.py registers a ResNet plugin there, which would apply to every later step.
+    """
+    from modelopt.torch.quantization.plugins import custom
+
+    plugins = custom.CUSTOM_POST_CONVERSION_PLUGINS
+    saved = set(plugins)
+    try:
+        yield
+    finally:
+        plugins.clear()
+        plugins.update(saved)
+
+
 def _preload(modules: Iterable[str]) -> None:
     """Import what ``_restored_state`` snapshots, so a first step cannot change it unnoticed."""
     for name in modules:
diff --git a/tests/examples/torch_onnx/conftest.py b/tests/examples/torch_onnx/conftest.py
new file mode 100644
index 00000000000..f226fbee9f9
--- /dev/null
+++ b/tests/examples/torch_onnx/conftest.py
@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run the torch_onnx steps in one long-lived worker; see ``_test_utils.examples.example_runner``."""
+
+import pytest
+from _test_utils.examples.example_runner import ExampleRunner, keep_post_conversion_plugins
+from _test_utils.examples.run_command import set_in_process_runner
+
+# hf_embedding_quant_to_onnx.py has not been checked for repeated runs in one process, so it keeps
+# running as a subprocess.
+_SCRIPTS = {"torch_quant_to_onnx.py"}
+
+
+@pytest.fixture(scope="session")
+def _example_runner():
+    runner = ExampleRunner(_SCRIPTS, hooks=(keep_post_conversion_plugins,))
+    yield runner
+    runner.close()
+
+
+@pytest.fixture(autouse=True)
+def _run_steps_in_the_worker(_example_runner):
+    """Install the runner per test, not per session.
+
+    The hook is a module global in ``run_command``, so leaving it installed would also reach any
+    other example suite collected later in the same session (``pytest tests/examples``).
+    """
+    set_in_process_runner(_example_runner)
+    try:
+        yield
+    finally:
+        set_in_process_runner(None)
diff --git a/tests/unit/test_example_runner.py b/tests/unit/test_example_runner.py
index 65d03dad3b4..47b369bd865 100644
--- a/tests/unit/test_example_runner.py
+++ b/tests/unit/test_example_runner.py
@@ -28,7 +28,12 @@
 
 import pytest
 from _test_utils.examples import run_command
-from _test_utils.examples.example_runner import ExampleRunner, StepFailed, StepTimeout
+from _test_utils.examples.example_runner import (
+    ExampleRunner,
+    StepFailed,
+    StepTimeout,
+    keep_post_conversion_plugins,
+)
 
 pytestmark = pytest.mark.usefixtures("skip_on_windows")
 
@@ -386,6 +391,23 @@ def hook():
     assert "ran" not in failure.value.output
 
 
+def test_keep_post_conversion_plugins_undoes_the_registrations_of_a_step(tmp_path):
+    from modelopt.torch.quantization.plugins.custom import CUSTOM_POST_CONVERSION_PLUGINS
+
+    before = set(CUSTOM_POST_CONVERSION_PLUGINS)
+    with _in_process_lane(tmp_path, (keep_post_conversion_plugins,)) as lane:
+        lane.write(
+            "s.py",
+            """
+            from modelopt.torch.quantization.plugins.custom import CUSTOM_POST_CONVERSION_PLUGINS
+
+            CUSTOM_POST_CONVERSION_PLUGINS.add(lambda model: None)
+            """,
+        )
+        lane.run("s.py")
+    assert before == CUSTOM_POST_CONVERSION_PLUGINS
+
+
 def test_worker_is_reused_until_a_step_fails(worker_lane):
     worker_lane.write("pid.py", "import os; print('pid', os.getpid())")
     worker_lane.write("bad.py", "raise SystemExit(1)")

From 846c0e4ade25ae3d415f9c7fd58a8f1957216706 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 09:49:47 +0000
Subject: [PATCH 32/34] Make test_example_runner.py collectable on Windows

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>
---
 tests/unit/test_example_runner.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_example_runner.py b/tests/unit/test_example_runner.py
index 47b369bd865..7660676fbbf 100644
--- a/tests/unit/test_example_runner.py
+++ b/tests/unit/test_example_runner.py
@@ -37,6 +37,10 @@
 
 pytestmark = pytest.mark.usefixtures("skip_on_windows")
 
+# Decorators run when the module is imported, which also happens on Windows (the tests then skip),
+# and signal has no SIGKILL there.
+_SIGKILL = getattr(signal, "SIGKILL", 9)
+
 
 class _Lane:
     """One example directory of a runner, with helpers to write and run its scripts."""
@@ -448,7 +452,7 @@ def test_timeout_kills_the_hung_step_with_its_children_and_the_next_step_gets_a_
 
 @pytest.mark.parametrize(
     ("death", "returncode"),
-    [("os._exit(7)", 7), ("os.kill(os.getpid(), signal.SIGKILL)", -signal.SIGKILL)],
+    [("os._exit(7)", 7), ("os.kill(os.getpid(), signal.SIGKILL)", -_SIGKILL)],
 )
 def test_worker_death_fails_the_step_not_the_session(worker_lane, death, returncode):
     worker_lane.write("die.py", f"import os, signal\nprint('about to die', flush=True)\n{death}")

From 1e2d732fcf5a5b8c16984fc5a0900bec5e7dfffe Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 09:49:49 +0000
Subject: [PATCH 33/34] Restore ConstraintsFunc._get_true_latency after
 test_search_constraints

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/nas/test_full_algorithms_vision.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/unit/torch/nas/test_full_algorithms_vision.py b/tests/unit/torch/nas/test_full_algorithms_vision.py
index f6c30c1da52..b3d81881696 100644
--- a/tests/unit/torch/nas/test_full_algorithms_vision.py
+++ b/tests/unit/torch/nas/test_full_algorithms_vision.py
@@ -117,7 +117,7 @@ def run_search():
         pytest.param(False, False, marks=pytest.mark.manual),
     ],
 )
-def test_search_constraints(flops_only: bool, bounded_latency: bool):
+def test_search_constraints(flops_only: bool, bounded_latency: bool, monkeypatch):
     def _fake_latency(self, model, precomputed=None):
         coeff_params = 0.00 if flops_only else 0.75
         if not hasattr(_fake_latency, "call_counter"):
@@ -130,8 +130,9 @@ def _fake_latency(self, model, precomputed=None):
     dummy_input = (*args, kwargs)
     model = mtn.convert(model, mode="autonas")
 
-    # get fake latency and spoof ConstraintsFunc class
-    ConstraintsFunc._get_true_latency = _fake_latency
+    # get fake latency and spoof ConstraintsFunc class (undone after the test: a class-level patch
+    # that stays behind changes the latency of every later test in the same process)
+    monkeypatch.setattr(ConstraintsFunc, "_get_true_latency", _fake_latency)
 
     # now setup constraints functor
     constraints_func = ConstraintsFunc(model, {}, dummy_input, {"device": None})

From 954ac755b20509ccd4875d2d8b7b6caa98bcabad Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 09:49:53 +0000
Subject: [PATCH 34/34] Give the unit tests 120 s per test on every platform

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>
---
 tests/conftest.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index af77546a244..b1685e3fe63 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -58,7 +58,10 @@ def pytest_addoption(parser):
     "gpu_trtllm": 60,
     "gpu_vllm": 60,
     "regression": 180,
-    "unit": 120 if platform.system() == "Windows" else 60,
+    # Tests that spawn processes take 2-3x longer under pytest-xdist: 4 workers and their children
+    # share the 4 vCPUs of a CI runner (test_dist.py::test_data_parallel: 8-15 s serial, 25-35 s with
+    # xdist, over 60 s on a slow runner).
+    "unit": 120,
 }