From 861c459a8a0ea70a89c39e561de78de3df011793 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 17 Apr 2026 12:11:05 -0700
Subject: [PATCH 1/9] Use nemo:26.02 container for megatron gpu tests

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/_example_tests_runner.yml   |  1 -
 .github/workflows/example_tests.yml           | 14 +++---
 .github/workflows/gpu_tests.yml               | 20 +++++---
 .../torch/peft/plugins/test_megatron_peft.py  | 21 ++------
 tox.ini                                       | 48 +++++++++----------
 5 files changed, 49 insertions(+), 55 deletions(-)

diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index b34aa87e0de..8adadbac7af 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -48,7 +48,6 @@ jobs:
       - name: Install dependencies
         run: |
           # use `python -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers
-          pip uninstall -y nvidia-modelopt
           python -m pip install ".${{ inputs.pip_install_extras }}"
 
           if [[ "${{ inputs.example }}" == *"diffusers"* ]]; then
diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 73f5c1fa9b7..620be113390 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -118,8 +118,8 @@ jobs:
       pip_install_extras: "[hf,dev-test]"
       runner: linux-amd64-gpu-rtxpro6000-latest-2
 
-  ##### NeMo Example Tests #####
-  nemo-pr:
+  ##### Megatron Example Tests #####
+  megatron-pr:
     needs: [check-file-changes, wait-checks]
     if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
     strategy: &nemo_strategy
@@ -135,7 +135,7 @@ jobs:
       pip_install_extras: "[hf,puzzletron,dev-test]"
       runner: linux-amd64-gpu-rtxpro6000-latest-1
 
-  nemo-non-pr:
+  megatron-non-pr:
     if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
     strategy: *nemo_strategy
     uses: ./.github/workflows/_example_tests_runner.yml
@@ -160,7 +160,7 @@ jobs:
     with:
       docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3"
       example: ${{ matrix.example }}
-      pip_install_extras: "[all,dev-test]"
+      pip_install_extras: "[onnx,hf,dev-test]"
       runner: linux-amd64-gpu-rtxpro6000-latest-1
 
   onnx-non-pr:
@@ -171,14 +171,14 @@ jobs:
     with:
       docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3"
       example: ${{ matrix.example }}
-      pip_install_extras: "[all,dev-test]"
+      pip_install_extras: "[onnx,hf,dev-test]"
       runner: linux-amd64-gpu-rtxpro6000-latest-2
 
   ##### Required Check for PR #####
   example-pr-required-check:
     # Run even if example tests are skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [check-file-changes, torch-pr, trtllm-pr, nemo-pr, onnx-pr]
+    needs: [check-file-changes, torch-pr, trtllm-pr, megatron-pr, onnx-pr]
     runs-on: ubuntu-latest
     steps:
       - name: Required GPU tests did not succeed
@@ -187,7 +187,7 @@ jobs:
           (needs.check-file-changes.outputs.any_changed == 'true' && (
             needs.torch-pr.result != 'success' ||
             needs.trtllm-pr.result != 'success' ||
-            needs.nemo-pr.result != 'success' ||
+            needs.megatron-pr.result != 'success' ||
             needs.onnx-pr.result != 'success'
           ))
         run: exit 1
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 30b47bb2160..d2b06229888 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -42,7 +42,9 @@ jobs:
             .github/workflows/gpu_tests.yml
             modelopt/**
             tests/gpu/**
+            tests/gpu_megatron/**
             tests/gpu_regression/**
+            tests/gpu_trtllm/**
             examples/speculative_decoding/**
             examples/dataset/**
             modelopt_recipes/general/speculative_decoding/**
@@ -71,13 +73,13 @@ jobs:
             timeout: 60
             container_image: pytorch:26.01-py3
             # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers until https://github.com/tbenthompson/cppimport/pull/98
-          - example: gpu-regression
+          - example: gpu_regression
             timeout: 15
             container_image: pytorch:26.01-py3
-          - example: gpu-megatron
+          - example: gpu_megatron
             timeout: 45
-            container_image: pytorch:26.01-py3
-          - example: gpu-trtllm
+            container_image: nemo:26.02
+          - example: gpu_trtllm
             timeout: 30
             container_image: tensorrt-llm/release:1.3.0rc10
     runs-on: linux-amd64-gpu-rtxpro6000-latest-1
@@ -99,8 +101,14 @@ jobs:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
         run: |
-          pip install tox-current-env
-          COV_ARGS="--cov" tox -e cuda13-${{ matrix.example }} --current-env
+          # nemo containers use uv venvs which is not compatible with tox-current-env, so run tests directly
+          if [[ "${{ matrix.example }}" == "gpu_megatron" ]]; then
+            python -m pip install -e .[hf,dev-test]
+            python -m pytest tests/gpu_megatron --cov
+          else
+            python -m pip install tox tox-current-env
+            COV_ARGS="--cov" python -m tox -e cuda13-${{ matrix.example }} --current-env
+          fi
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
diff --git a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py
index 4f76b3f3ae8..50e3f43ab7d 100644
--- a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py
+++ b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py
@@ -20,8 +20,11 @@
 import torch
 import torch.nn.init as init
 from _test_utils.torch.megatron.models import get_mcore_gpt_model
-from _test_utils.torch.megatron.utils import initialize_for_megatron
-from megatron.core import dist_checkpointing
+from _test_utils.torch.megatron.utils import (
+    initialize_for_megatron,
+    load_distributed_checkpoint,
+    save_distributed_checkpoint,
+)
 
 import modelopt.torch.peft as mtpeft
 import modelopt.torch.quantization as mtq
@@ -148,20 +151,6 @@
 }
 
 
-def save_distributed_checkpoint(checkpoint_path, gpt_model):
-    sharded_state_dict = gpt_model.sharded_state_dict(prefix="")
-    dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
-
-
-def load_distributed_checkpoint(checkpoint_path, gpt_model):
-    sharded_state_dict = gpt_model.sharded_state_dict(prefix="")
-    checkpoint = dist_checkpointing.load(
-        sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path
-    )
-    gpt_model.load_state_dict(checkpoint)
-    return gpt_model
-
-
 def _gpt_model_provider(tp_size: int, hidden_size=256, vocab_size=64, meta_device=False):
     """Build the model."""
 
diff --git a/tox.ini b/tox.ini
index 7bfa1e41e57..ffbc95227cd 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,10 +1,9 @@
 [tox]
 envlist=
     pre-commit-all
-    py312-torch210-tf_latest-unit
+    py312-torch211-tf_latest-unit
     cuda13-gpu
-    cuda13-gpu-regression
-    cuda13-gpu-megatron
+    cuda13-gpu_regression
 skipsdist = True
 toxworkdir = /tmp/{env:USER}-modelopt-tox
 passenv =
@@ -59,44 +58,43 @@ commands =
 ###########################################################
 # GPU test environments (Should be used with --current-env)
 ###########################################################
+# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
 [testenv:cuda13-gpu]
 commands_pre =
     # Install deps here so that it gets installed even in --current-env
-    pip install --no-build-isolation git+https://github.com/Dao-AILab/fast-hadamard-transform.git
-    pip install -e .[all,dev-test]
+    python -m pip install --no-build-isolation git+https://github.com/Dao-AILab/fast-hadamard-transform.git
+    python -m pip install -e .[all,dev-test]
 
     # Install cupy-cuda13x for INT4 ONNX quantization (default is cupy-cuda12x)
-    pip uninstall -y cupy-cuda12x
-    pip install cupy-cuda13x
+    python -m pip uninstall -y cupy-cuda12x
+    python -m pip install cupy-cuda13x
 
     # Install mamba and causal-conv1d for Nemotron tests
-    pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git
-    pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git
+    python -m pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git
+    python -m pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git
 commands =
     python -m pytest tests/gpu {env:COV_ARGS:}
 
-[testenv:cuda13-gpu-regression]
+[testenv:cuda13-gpu_regression]
 commands_pre =
-    pip install -e .[hf,dev-test]
+    python -m pip install -e .[hf,dev-test]
 commands =
     python -m pytest tests/gpu_regression {env:COV_ARGS:}
 
-[testenv:cuda13-gpu-megatron]
+# Container: nvcr.io/nvidia/nemo:26.02 or later
+# NOTE: tox is bypassed for this env in CI (see gpu_tests.yml) because tox-current-env is
+# incompatible with uv venvs, and any new tox env would lack NeMo/Megatron packages from /opt/venv.
+# [testenv:cuda13-gpu_megatron]
+# commands_pre =
+#     python -m pip install -e .[hf,dev-test]
+# commands =
+#     python -m pytest tests/gpu_megatron {env:COV_ARGS:}
+
+# Container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 or later
+[testenv:cuda13-gpu_trtllm]
 commands_pre =
     # Install deps here so that it gets installed even in --current-env
-    # Temporarily disable latest mcore until we fix its nvidia-resiliency-ext dependency
-    pip install 'megatron-core<0.17.0'
-    pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git
-    pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git
-    pip install -e .[hf,dev-test]
-commands =
-    python -m pytest tests/gpu_megatron {env:COV_ARGS:}
-
-[testenv:cuda13-gpu-trtllm]
-# Expected to be run in TRT-LLM container
-commands_pre =
-    # Install deps here so that it gets installed even in --current-env
-    pip install -e .[hf,dev-test]
+    python -m pip install -e .[hf,dev-test]
 commands =
     python -m pytest tests/gpu_trtllm {env:COV_ARGS:}
 

From 84408a23ae2c7c3378abf9c843ef8e9c7d5826ae Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 11:34:30 +0530
Subject: [PATCH 2/9] Replace tox with nox for all CI test sessions (#1290)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Type of change: New feature / infrastructure improvement

Replaces `tox` + `tox-current-env` with `nox` for all test, lint, docs,
and wheel build sessions. The primary motivation was that
`tox-current-env` is incompatible with uv venvs in NGC containers (e.g.
NeMo's `/opt/venv`) — it picks the system Python via
`sys._base_executable` instead of the container's venv Python which has
megatron packages pre-installed.

Key changes:
- **`noxfile.py`** replaces `tox.ini` with GPU, CPU unit,
partial-install, pre-commit, docs, and wheel sessions
- **GPU sessions** use `venv_backend="none"` (run directly in container
env) and `python -m pip/pytest` to avoid PATH mismatches
- **CPU unit sessions** use 2-level `@nox.parametrize` over torch ×
transformers versions — any combination is selectable e.g. `nox -s
"unit-3.12(torch_211-tf_latest)"`
- **uv** is set as the default venv backend for CPU sessions (faster
installs); `envdir=/tmp/.nox` avoids permission errors in mounted
container directories
- All CI workflows updated to use `pip install nox uv && nox -s
<session>`

```bash
pip install nox uv                              # install once
nox -l                                          # list all sessions
nox -s "unit-3.12(torch_211-tf_latest)"         # default unit tests
nox -s "unit-3.12(torch_28-tf_min)"             # torch 2.8 + min transformers
nox -s gpu_megatron                             # run inside NeMo container
```

- Ran `nox -l` to verify all session names
- Ran `gpu_megatron` session locally inside NeMo container — confirmed
it uses `/opt/venv/bin/python` correctly

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: N/A — CI infrastructure only
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ (added `nox`
and `uv` to `dev-test`, both Apache-2.0)
- Did you write any new necessary tests?: N/A
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A — no user-facing changes

Supersedes the tox-current-env workaround in the parent branch.

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .github/CODEOWNERS                 |   2 +-
 .github/workflows/code_quality.yml |   2 +-
 .github/workflows/gpu_tests.yml    |  16 +--
 .github/workflows/pages.yml        |   4 +-
 .github/workflows/release.yml      |   6 +-
 .github/workflows/unit_tests.yml   |  20 +--
 .vscode/settings.json              |   1 -
 CLAUDE.md                          |   8 +-
 CONTRIBUTING.md                    |   2 +-
 noxfile.py                         | 193 +++++++++++++++++++++++++++++
 pyproject.toml                     |  13 +-
 tools/launcher/tests/conftest.py   |   4 +-
 tox.ini                            | 147 ----------------------
 13 files changed, 229 insertions(+), 189 deletions(-)
 create mode 100644 noxfile.py
 delete mode 100644 tox.ini

diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 9bac2756b8b..19188206ea0 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -9,7 +9,7 @@ LICENSE @NVIDIA/modelopt-setup-codeowners
 LICENSE_HEADER @NVIDIA/modelopt-setup-codeowners
 pyproject.toml @NVIDIA/modelopt-setup-codeowners
 SECURITY.md @NVIDIA/modelopt-setup-codeowners
-tox.ini @NVIDIA/modelopt-setup-codeowners
+noxfile.py @NVIDIA/modelopt-setup-codeowners
 uv.lock @NVIDIA/modelopt-setup-codeowners
 
 # Library
diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml
index 723f3af1209..cd877d18d24 100644
--- a/.github/workflows/code_quality.yml
+++ b/.github/workflows/code_quality.yml
@@ -24,4 +24,4 @@ jobs:
         with:
           extra_args: --results=verified,unknown
       - name: Run code quality checks
-        run: pip install tox && tox -e pre-commit-all
+        run: pip install nox uv && nox -s pre_commit_all
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index d2b06229888..95d633ec7c7 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -41,16 +41,16 @@ jobs:
           files: |
             .github/workflows/gpu_tests.yml
             modelopt/**
+            noxfile.py
+            pyproject.toml
             tests/gpu/**
             tests/gpu_megatron/**
-            tests/gpu_regression/**
             tests/gpu_trtllm/**
+            tests/gpu_regression/**
             examples/speculative_decoding/**
             examples/dataset/**
             modelopt_recipes/general/speculative_decoding/**
             tools/launcher/**
-            pyproject.toml
-            tox.ini
           fail_on_initial_diff_error: true
   wait-checks:
     needs: [check-file-changes]
@@ -101,14 +101,8 @@ jobs:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
         run: |
-          # nemo containers use uv venvs which is not compatible with tox-current-env, so run tests directly
-          if [[ "${{ matrix.example }}" == "gpu_megatron" ]]; then
-            python -m pip install -e .[hf,dev-test]
-            python -m pytest tests/gpu_megatron --cov
-          else
-            python -m pip install tox tox-current-env
-            COV_ARGS="--cov" python -m tox -e cuda13-${{ matrix.example }} --current-env
-          fi
+          pip install nox
+          nox -s ${{ matrix.example }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml
index 6789142f97f..eda19778594 100644
--- a/.github/workflows/pages.yml
+++ b/.github/workflows/pages.yml
@@ -27,7 +27,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Build docs
-        run: pip install tox && tox -e build-docs
+        run: pip install nox uv && nox -s docs
       - name: Upload docs artifact
         if: github.event_name == 'push' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'
         uses: actions/upload-artifact@v4
@@ -47,7 +47,7 @@ jobs:
       - uses: ./.github/actions/ubuntu-setup
       - name: Build docs
         if: github.event.action != 'closed'
-        run: pip install tox && tox -e build-docs
+        run: pip install nox uv && nox -s docs
       - name: Deploy / remove PR preview
         uses: rossjrw/pr-preview-action@v1
         with:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index e3e62d78def..35edb9e53ef 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -13,13 +13,11 @@ jobs:
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
-      - name: Install dependencies
-        run: pip install tox
       - name: Run basic unit tests
-        run: tox -e py312-torch29-tf_latest-unit
+        run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)"
       - name: Build Wheel
         run: |
-          tox -e build-wheel
+          nox -s build_wheel
           echo "WHEEL_PATH=$(find dist -name "*.whl" | head -n 1)" >> $GITHUB_ENV
       - name: Upload GitHub Release Artifact
         env:
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index afa57fc3a86..29bd48ebd07 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -9,9 +9,9 @@ on:
     paths:
       - ".github/workflows/unit_tests.yml"
       - "modelopt/**"
-      - "tests/unit/**"
+      - "noxfile.py"
       - "pyproject.toml"
-      - "tox.ini"
+      - "tests/unit/**"
       - "tools/launcher/**"
   schedule:
     - cron: "0 0 * * *" # Nightly
@@ -38,7 +38,10 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Run unit tests
-        run: pip install tox && COV_ARGS="--cov" tox -e py312-torch211-tf_latest-unit
+        env:
+          COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
+          COVERAGE_FILE: ${{ github.workspace }}/.coverage
+        run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)"
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v5
         with:
@@ -57,8 +60,7 @@ jobs:
         with:
           python-version: "3.12"
       - name: Run unit tests (without coverage)
-        # Some issues with torch 2.10 on Windows, so using 2.9 for now
-        run: pip install tox && tox -e py312-torch29-tf_latest-unit
+        run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)"
   multi-py:
     if: github.event_name == 'pull_request'
     needs: [linux]
@@ -74,7 +76,7 @@ jobs:
         with:
           python-version: "3.${{ matrix.py }}"
       - name: Run unit tests
-        run: pip install tox && tox -e py3${{ matrix.py }}-torch211-tf_latest-unit
+        run: pip install nox uv && nox -s "unit-3.${{ matrix.py }}(torch_211, tf_latest)"
   multi-torch:
     if: github.event_name == 'pull_request'
     needs: [linux]
@@ -88,7 +90,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Run unit tests
-        run: pip install tox && tox -e py312-torch${{ matrix.torch }}-tf_latest-unit
+        run: pip install nox uv && nox -s "unit-3.12(torch_${{ matrix.torch }}, tf_latest)"
   multi-transformers:
     if: github.event_name == 'pull_request'
     needs: [linux]
@@ -102,7 +104,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Run unit tests
-        run: pip install tox && tox -e py312-torch211-tf_${{ matrix.tf }}-unit
+        run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_min)"
   launcher:
     if: github.event_name == 'pull_request'
     needs: [linux]
@@ -133,7 +135,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Run unit tests
-        run: pip install tox && tox -e py312-partial-unit-${{ matrix.test-env }}
+        run: pip install nox uv && nox -s "partial_unit(subset='${{ matrix.test-env }}')"
   unit-pr-required-check:
     # Run even if some jobs are skipped
     if: ${{ github.event_name == 'pull_request' && always() }}
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 0e8465ad38d..5a1c278b0f7 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -19,7 +19,6 @@
         ".mypy_cache": true,
         ".pytest_cache": true,
         ".ruff_cache": true,
-        ".tox": true,
         "**/__pycache__/**": true,
         "**/*.pyc": true,
         "**/runs": true,
diff --git a/CLAUDE.md b/CLAUDE.md
index d07775f3a85..4af38586788 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -37,9 +37,9 @@ Primarily Python codebase with optional C++/CUDA extensions supporting PyTorch,
 | Pattern match | `pytest tests/unit -k "test_quantize"` |
 | Lint + format (all files) | `pre-commit run --all-files` |
 | Lint (diff only) | `pre-commit run --from-ref origin/main --to-ref HEAD` |
-| Run via tox (CPU unit) | `tox -e py312-torch210-tf_latest-unit` |
-| Build docs | `tox -e build-docs` |
-| Build wheel | `tox -e build-wheel` |
+| Run via nox (CPU unit) | `nox -s "unit-3.12(torch_211, tf_latest)"` |
+| Build docs | `nox -s docs` |
+| Build wheel | `nox -s build_wheel` |
 
 ## Architecture
 
@@ -104,7 +104,7 @@ A **recipe** is a declarative YAML specification of an optimization configuratio
 | `modelopt_recipes/general/ptq/` | Built-in PTQ recipe YAML files |
 | `pyproject.toml` | Optional dependency groups (`[onnx]`, `[hf]`, `[all]`, `[dev]`); ruff, mypy, pytest, bandit, and coverage config |
 | `.pre-commit-config.yaml` | Pre-commit hooks (ruff, mypy, clang-format, license headers) |
-| `tox.ini` | Test environment definitions |
+| `noxfile.py` | Test session definitions |
 
 ## Design Patterns
 
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 23f07850a6b..53e879a7bf3 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -89,7 +89,7 @@ We use [pytest](https://docs.pytest.org/) for all tests. For any new features /
 - `tests/gpu_trtllm`: Fast GPU-based unit tests for the core ModelOpt library for TensorRT-LLM features. In most cases, they should not take more than a few seconds to run.
 - `tests/examples`: Integration tests for ModelOpt examples. They should not take more than a few minutes to run. Please refer to [example test README](./tests/examples/README.md) for more details.
 
-Please refer to [tox.ini](./tox.ini) for more details on how to run the tests and their dependencies.
+Please refer to [noxfile.py](./noxfile.py) for more details on how to run the tests and their dependencies.
 
 ## ✍️ Signing your work
 
diff --git a/noxfile.py b/noxfile.py
new file mode 100644
index 00000000000..3d6d9d66eca
--- /dev/null
+++ b/noxfile.py
@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+"""Nox session definitions for testing, linting, docs, and wheel builds.
+
+Usage:
+    python -m pip install nox uv                                                    # install nox and uv (once)
+    nox -l                                                                          # list all sessions
+    nox -s gpu_megatron                                                             # run a GPU session (inside container)
+    nox -s "unit-3.12(torch_211, tf_latest)"                                        # run a specific unit test combination
+    nox -s "unit-3.12(torch_211, tf_latest)" -R                                     # force-recreate venv (e.g. after dep changes)
+    COVERAGE_PROCESS_START=pyproject.toml nox -s "unit-3.12(torch_211, tf_latest)"  # with coverage
+"""
+
+import glob
+import os
+import shutil
+
+import nox
+
+nox.options.default_venv_backend = "uv" if shutil.which("uv") else "virtualenv"
+nox.options.envdir = "/tmp/.nox"
+nox.options.reuse_existing_virtualenvs = True
+
+TORCH_VERSIONS = {
+    "torch_28": "torchvision~=0.23.0",
+    "torch_29": "torchvision~=0.24.0",
+    "torch_210": "torchvision~=0.25.0",
+    "torch_211": "torchvision~=0.26.0",
+}
+
+TRANSFORMERS_VERSIONS = {
+    "tf_latest": None,
+    "tf_min": "transformers~=4.56.0",
+}
+
+
+def _cov_args():
+    """Return --cov when COVERAGE_PROCESS_START is set (CI only)."""
+    return ["--cov"] if os.environ.get("COVERAGE_PROCESS_START") else []
+
+
+# ─── CPU unit tests ───────────────────────────────────────────────────────────
+@nox.session(python=["3.10", "3.11", "3.12", "3.13"])
+@nox.parametrize("tf_ver", [nox.param(k, id=k) for k in TRANSFORMERS_VERSIONS])
+@nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS])
+def unit(session, torch_ver, tf_ver):
+    """Unit tests — parametrized over torch and transformers versions."""
+    session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]")
+    tf_pin = TRANSFORMERS_VERSIONS[tf_ver]
+    if tf_pin:
+        session.install(tf_pin)
+    session.run("python", "-m", "pytest", "tests/unit", *_cov_args())
+
+
+@nox.session(python="3.12")
+@nox.parametrize("subset", ["onnx", "torch", "torch_deploy"])
+def partial_unit(session, subset):
+    """Unit tests with partial installs."""
+    if subset == "onnx":
+        session.install("torchvision~=0.26.0", ".[onnx,dev-test]")
+        session.run("python", "-m", "pytest", "tests/unit/onnx")
+    elif subset == "torch":
+        session.install("megatron-core", ".[dev-test]")
+        session.run(
+            "python",
+            "-m",
+            "pytest",
+            "tests/unit/torch",
+            "--ignore=tests/unit/torch/deploy",
+            "--ignore=tests/unit/torch/puzzletron",
+        )
+    else:  # torch_deploy
+        session.install(".[onnx,dev-test]")
+        session.run("python", "-m", "pytest", "tests/unit/torch/deploy")
+
+
+# ─── GPU sessions (run inside containers — no new venv) ──────────────────────
+# `venv_backend="none"` skips creating a new venv so the session runs directly in the container's
+# 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).
+# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
+@nox.session(venv_backend="none")
+def gpu(session):
+    # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers
+    # until https://github.com/tbenthompson/cppimport/pull/98
+    session.run(
+        "python",
+        "-m",
+        "pip",
+        "install",
+        "--no-build-isolation",
+        "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+    )
+    session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test]")
+    session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x")
+    session.run("python", "-m", "pip", "install", "cupy-cuda13x")
+    session.run(
+        "python",
+        "-m",
+        "pip",
+        "install",
+        "--no-build-isolation",
+        "git+https://github.com/state-spaces/mamba.git",
+        "git+https://github.com/Dao-AILab/causal-conv1d.git",
+    )
+    session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
+
+
+# Container: nvcr.io/nvidia/nemo:26.02 or later
+@nox.session(venv_backend="none")
+def gpu_megatron(session):
+    session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
+    session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args())
+
+
+# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
+@nox.session(venv_backend="none")
+def gpu_regression(session):
+    session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
+    session.run("python", "-m", "pytest", "tests/gpu_regression", *_cov_args())
+
+
+# Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 or later
+@nox.session(venv_backend="none")
+def gpu_trtllm(session):
+    session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
+    session.run("python", "-m", "pytest", "tests/gpu_trtllm", *_cov_args())
+
+
+# ─── Code quality ─────────────────────────────────────────────────────────────
+@nox.session
+def pre_commit_all(session):
+    session.install("-e", ".[all,dev-lint]")
+    session.run("pre-commit", "run", "--all-files", "--show-diff-on-failure")
+
+
+@nox.session
+def pre_commit_diff(session):
+    session.install("-e", ".[all,dev-lint]")
+    session.run("pre-commit", "run", "--from-ref", "origin/main", "--to-ref", "HEAD")
+
+
+# ─── Docs ─────────────────────────────────────────────────────────────────────
+@nox.session
+def docs(session):
+    session.install("-e", ".[all,dev-docs]")
+    shutil.rmtree("docs/build", ignore_errors=True)
+    shutil.rmtree("docs/source/reference/generated", ignore_errors=True)
+    with session.chdir("docs"):
+        session.run(
+            "sphinx-build",
+            "source",
+            "build/html",
+            "--fail-on-warning",
+            "--show-traceback",
+            "--keep-going",
+        )
+
+
+@nox.session
+def docs_debug(session):
+    session.install("-e", ".[all,dev-docs]")
+    shutil.rmtree("docs/build", ignore_errors=True)
+    shutil.rmtree("docs/source/reference/generated", ignore_errors=True)
+    with session.chdir("docs"):
+        session.run("sphinx-autobuild", "source", "build/html", "--host", "0.0.0.0")
+
+
+# ─── Wheel build ──────────────────────────────────────────────────────────────
+@nox.session
+def build_wheel(session):
+    shutil.rmtree("build", ignore_errors=True)
+    session.install("twine")
+    session.run("pip", "wheel", "--no-deps", "--wheel-dir=dist", ".")
+    wheels = glob.glob("dist/*.whl")
+    session.run("twine", "check", *wheels)
+    (modelopt_wheel,) = glob.glob("dist/nvidia_modelopt-*.whl")
+    session.install(modelopt_wheel, "-f", "dist")
+    with session.chdir("dist"):
+        session.run("python", "-c", "import modelopt; print(modelopt.__version__)")
diff --git a/pyproject.toml b/pyproject.toml
index 25cb6338aa5..100f149ffe5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -82,7 +82,7 @@ hf = [
     "peft>=0.17.0",
     "sentencepiece>=0.2.1",                                                           # Also implicitly used in test_unified_export_megatron, test_vllm_fakequant_megatron_export
     "tiktoken",
-    "transformers>=4.56",                                                             # Should match modelopt/torch/__init__.py and tox.ini
+    "transformers>=4.56",                                                             # Should match modelopt/torch/__init__.py and noxfile.py
     "wonderwords",
 ]
 
@@ -111,17 +111,17 @@ dev-docs = [
     "sphinx-togglebutton>=0.3.2",
 ]
 dev-test = [
-    "coverage[toml]>=7.13.0",  # a1_coverage.pth for subprocess tracking requires this
+    "coverage[toml]>=7.13.0", # a1_coverage.pth for subprocess tracking requires this
     "pytest",
     "pytest-cov",
     "pytest-instafail",
     "pytest-timeout",
     "timm",
-    "torchprofile==0.0.4",     # optional dependency for modelopt.torch
+    "torchprofile==0.0.4",    # optional dependency for modelopt.torch
     "torchvision",
     "torch-geometric",
-    "tox>4.18",
-    "tox-current-env>=0.0.12",
+    "uv",
+    "nox",
 ]
 # Compound extras via self-references
 all = ["nvidia-modelopt[onnx,hf,puzzletron]"]
@@ -208,6 +208,7 @@ extend-ignore = [
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401", "F403"]
 "examples/*" = ["D"]
+"noxfile.py" = ["D", "E501"]
 "tests/*" = ["B017", "D", "E402", "PT012"]
 "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py)
 "*.ipynb" = [
@@ -279,7 +280,7 @@ module = ["examples.*"]
 disable_error_code = ["attr-defined"]
 
 [tool.bandit]
-exclude_dirs = [".github/", "examples/", "tests/"]
+exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"]
 # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config.
 # Use of `# nosec BXXX` requires special approval
 skips = [
diff --git a/tools/launcher/tests/conftest.py b/tools/launcher/tests/conftest.py
index da98e8ad6cb..6d6cb56f635 100644
--- a/tools/launcher/tests/conftest.py
+++ b/tools/launcher/tests/conftest.py
@@ -20,8 +20,8 @@
     uv pip install pytest
     uv run python3 -m pytest tests/ -v
 
-Or via tox from Model-Optimizer root:
-    tox -e py312-launcher
+Or via nox from Model-Optimizer root:
+    nox -s "unit-3.12(torch_211-tf_latest)"
 """
 
 import os
diff --git a/tox.ini b/tox.ini
deleted file mode 100644
index ffbc95227cd..00000000000
--- a/tox.ini
+++ /dev/null
@@ -1,147 +0,0 @@
-[tox]
-envlist=
-    pre-commit-all
-    py312-torch211-tf_latest-unit
-    cuda13-gpu
-    cuda13-gpu_regression
-skipsdist = True
-toxworkdir = /tmp/{env:USER}-modelopt-tox
-passenv =
-    SETUPTOOLS_SCM_PRETEND_VERSION
-
-############################
-# CPU Unit test environments
-############################
-[testenv:{py310,py311,py312,py313}-torch{28,29,210,211}-tf_{min,latest}-unit]
-deps =
-    # torch version auto-selected based on torchvision version
-    torch28: torchvision~=0.23.0
-    torch29: torchvision~=0.24.0
-    torch210: torchvision~=0.25.0
-    torch211: torchvision~=0.26.0
-
-    -e .[all,dev-test]
-
-    # Should match pyproject.toml
-    tf_min: transformers~=4.56.0
-commands =
-    python -m pytest tests/unit {env:COV_ARGS:}
-
-
-#####################################################################
-# Environment to run unit tests with subset of dependencies installed
-#####################################################################
-[testenv:{py310,py311,py312,py313}-partial-unit-{onnx,torch,torch_deploy}]
-allowlist_externals =
-    bash, rm
-deps =
-    # Make sure torch 2.10 is used
-    torchvision~=0.26.0
-
-    # ONNX unit tests heavily rely on torch / torchvision
-    onnx: .[onnx,dev-test]
-    onnx: torchvision
-
-    # Install megatron-core to test torch-only install can still import plugins
-    torch: megatron-core
-    # diffusers is needed for unit tests of the sparse-attention/quantization diffusers backend
-    torch: diffusers
-    torch: .[dev-test]
-
-    torch_deploy: .[onnx,torch,dev-test]
-commands =
-    onnx: python -m pytest tests/unit/onnx
-    torch: python -m pytest tests/unit/torch --ignore tests/unit/torch/deploy --ignore tests/unit/torch/puzzletron
-    torch_deploy: python -m pytest tests/unit/torch/deploy
-
-
-###########################################################
-# GPU test environments (Should be used with --current-env)
-###########################################################
-# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
-[testenv:cuda13-gpu]
-commands_pre =
-    # Install deps here so that it gets installed even in --current-env
-    python -m pip install --no-build-isolation git+https://github.com/Dao-AILab/fast-hadamard-transform.git
-    python -m pip install -e .[all,dev-test]
-
-    # Install cupy-cuda13x for INT4 ONNX quantization (default is cupy-cuda12x)
-    python -m pip uninstall -y cupy-cuda12x
-    python -m pip install cupy-cuda13x
-
-    # Install mamba and causal-conv1d for Nemotron tests
-    python -m pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git
-    python -m pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git
-commands =
-    python -m pytest tests/gpu {env:COV_ARGS:}
-
-[testenv:cuda13-gpu_regression]
-commands_pre =
-    python -m pip install -e .[hf,dev-test]
-commands =
-    python -m pytest tests/gpu_regression {env:COV_ARGS:}
-
-# Container: nvcr.io/nvidia/nemo:26.02 or later
-# NOTE: tox is bypassed for this env in CI (see gpu_tests.yml) because tox-current-env is
-# incompatible with uv venvs, and any new tox env would lack NeMo/Megatron packages from /opt/venv.
-# [testenv:cuda13-gpu_megatron]
-# commands_pre =
-#     python -m pip install -e .[hf,dev-test]
-# commands =
-#     python -m pytest tests/gpu_megatron {env:COV_ARGS:}
-
-# Container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 or later
-[testenv:cuda13-gpu_trtllm]
-commands_pre =
-    # Install deps here so that it gets installed even in --current-env
-    python -m pip install -e .[hf,dev-test]
-commands =
-    python -m pytest tests/gpu_trtllm {env:COV_ARGS:}
-
-#############################################
-# Code quality checks on all files or on diff
-#############################################
-[testenv:{pre-commit}-{all,diff}]
-deps =
-    -e .[all,dev-lint]
-commands =
-    all: pre-commit run --all-files --show-diff-on-failure {posargs}
-    diff: pre-commit run --from-ref origin/main --to-ref HEAD {posargs}
-
-
-#########################
-# Run documentation build
-#########################
-[testenv:{build,debug}-docs]
-allowlist_externals =
-    rm
-deps =
-    -e .[all,dev-docs]
-changedir = docs
-commands_pre =
-    rm -rf build
-    rm -rf source/reference/generated
-commands =
-    sphinx-build source build/html --fail-on-warning --show-traceback --keep-going
-    debug: sphinx-autobuild source build/html --host 0.0.0.0
-
-
-#################
-# Run wheel build
-#################
-[testenv:build-wheel]
-allowlist_externals =
-    bash, cd, rm
-deps =
-    twine
-commands =
-    # Clean build directory to avoid any stale files getting into the wheel
-    rm -rf build
-
-    # Build and check wheel
-    pip wheel --no-deps --wheel-dir=dist .
-    twine check dist/*
-
-    # Install and test the wheel
-    bash -c "find dist -name 'nvidia_modelopt-*.whl' | xargs pip install -f dist"
-    bash -c "cd dist; python -c 'import modelopt; print(modelopt.__version__);'"

From a8f5667c72dc4296cfa747ae4548ea7e6a637f8b Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 17 Apr 2026 23:55:24 -0700
Subject: [PATCH 3/9] minor

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml      |  6 ++++--
 .github/workflows/gpu_tests.yml          |  9 +++++----
 .github/workflows/unit_tests.yml         |  6 ++++--
 tests/unit/torch/speculative/conftest.py | 21 +++++++++++++++++++++
 4 files changed, 34 insertions(+), 8 deletions(-)
 create mode 100644 tests/unit/torch/speculative/conftest.py

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 620be113390..cbccfaa6149 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -6,10 +6,12 @@ on:
     # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used
   schedule:
     - cron: "0 0 * * *" # Nightly
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
+
 
-# Cancel previous runs if new commit is pushed to the same PR
 concurrency:
+  # Cancel previous runs if new commit is pushed to the same PR
   group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }}
   cancel-in-progress: true
 
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 95d633ec7c7..fe86eaea952 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -6,10 +6,12 @@ on:
     # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used
   schedule:
     - cron: "0 0 * * *" # Nightly
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
+
 
-# Cancel previous runs if new commit is pushed to the same PR
 concurrency:
+  # Cancel previous runs if new commit is pushed to the same PR
   group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }}
   cancel-in-progress: true
 
@@ -101,8 +103,7 @@ jobs:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
         run: |
-          pip install nox
-          nox -s ${{ matrix.example }}
+          python -m pip install nox && nox -s ${{ matrix.example }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index 29bd48ebd07..8401cb9120e 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -15,10 +15,12 @@ on:
       - "tools/launcher/**"
   schedule:
     - cron: "0 0 * * *" # Nightly
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
+
 
-# Cancel previous runs if new commit is pushed
 concurrency:
+  # Cancel previous runs if new commit is pushed
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
   cancel-in-progress: true
 
diff --git a/tests/unit/torch/speculative/conftest.py b/tests/unit/torch/speculative/conftest.py
new file mode 100644
index 00000000000..083769483a3
--- /dev/null
+++ b/tests/unit/torch/speculative/conftest.py
@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+import platform
+
+
+# `Compiler: cl is not found` on Windows
+def pytest_ignore_collect(collection_path, config):
+    return platform.system() == "Windows"

From 6a070f71c2d509f1772fa996a12fdd86fabff5f5 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 02:50:35 -0700
Subject: [PATCH 4/9] Use nemo:26.04 for gpu_megatron tests and fix
 GPTModelExporter.save_pretrained race condition hang

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml                            | 2 +-
 modelopt/torch/export/unified_export_megatron.py           | 6 ++++--
 noxfile.py                                                 | 4 +++-
 pyproject.toml                                             | 7 ++++---
 .../torch/export/test_unified_export_megatron.py           | 3 +--
 tools/launcher/tests/conftest.py                           | 2 +-
 6 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index fe86eaea952..ed4fbe9de07 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -80,7 +80,7 @@ jobs:
             container_image: pytorch:26.01-py3
           - example: gpu_megatron
             timeout: 45
-            container_image: nemo:26.02
+            container_image: nemo:26.04
           - example: gpu_trtllm
             timeout: 30
             container_image: tensorrt-llm/release:1.3.0rc10
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index c901a2159d8..89b718623da 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -76,6 +76,7 @@
     from megatron.core.parallel_state import (
         get_pipeline_model_parallel_rank,
         get_pipeline_model_parallel_world_size,
+        get_tensor_model_parallel_rank,
     )
     from megatron.core.ssm.mamba_layer import MambaLayer
     from megatron.core.transformer.identity_op import IdentityOp
@@ -258,13 +259,14 @@ def save_pretrained(
         """
         pp_rank = get_pipeline_model_parallel_rank()
         pp_size = get_pipeline_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
 
         # We use the 1st PP rank to handle VLM because vision_models
         # and vision_proj only exist in the first stage.
-        is_first_stage_main_rank = pp_rank == 0
+        is_first_stage_main_rank = pp_rank == 0 and tp_rank == 0
         # We use the last PP rank to write the config because
         # medusa_heads and eagle_module only exist in the last stage.
-        is_last_stage_main_rank = pp_rank == pp_size - 1
+        is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0
 
         # Main export process
         layer_state_dicts = self.layer_state_dicts
diff --git a/noxfile.py b/noxfile.py
index 3d6d9d66eca..5cd8f8b5ed3 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -119,9 +119,11 @@ def gpu(session):
     session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
 
 
-# Container: nvcr.io/nvidia/nemo:26.02 or later
+# Container: nvcr.io/nvidia/nemo:26.04 or later
 @nox.session(venv_backend="none")
 def gpu_megatron(session):
+    # nemo:26.04 has transformers 5.x but tensorrt_llm 1.2.0 which does not support it causing import errors
+    session.run("python", "-m", "pip", "uninstall", "-y", "tensorrt_llm")
     session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
     session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args())
 
diff --git a/pyproject.toml b/pyproject.toml
index 100f149ffe5..fdd60b5193a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -112,16 +112,17 @@ dev-docs = [
 ]
 dev-test = [
     "coverage[toml]>=7.13.0", # a1_coverage.pth for subprocess tracking requires this
+    "nox",
     "pytest",
     "pytest-cov",
     "pytest-instafail",
     "pytest-timeout",
+    "uv",
+    # test-specific dependencies
     "timm",
-    "torchprofile==0.0.4",    # optional dependency for modelopt.torch
+    "torchprofile==0.0.4", # optional dependency for modelopt.torch
     "torchvision",
     "torch-geometric",
-    "uv",
-    "nox",
 ]
 # Compound extras via self-references
 all = ["nvidia-modelopt[onnx,hf,puzzletron]"]
diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
index 925c0e3c51c..8ccb3d42906 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -24,6 +24,7 @@
 from _test_utils.torch.megatron.models import get_mcore_gpt_model
 from _test_utils.torch.megatron.utils import get_forward
 from _test_utils.torch.transformers_models import create_tiny_llama_dir
+from safetensors import safe_open
 from safetensors.torch import save_file
 
 import modelopt.torch.quantization as mtq
@@ -275,8 +276,6 @@ def _test_qkv_slicing_gqa_tp2(tmp_path, rank, size):
 
     # Verify Q/K/V projections were exported (collect keys from all shard files)
     if rank == 0:
-        from safetensors import safe_open
-
         safetensors_files = list(export_dir.glob("*.safetensors"))
         assert safetensors_files, "no safetensors files found in export dir"
         keys = []
diff --git a/tools/launcher/tests/conftest.py b/tools/launcher/tests/conftest.py
index 6d6cb56f635..072518cc795 100644
--- a/tools/launcher/tests/conftest.py
+++ b/tools/launcher/tests/conftest.py
@@ -21,7 +21,7 @@
     uv run python3 -m pytest tests/ -v
 
 Or via nox from Model-Optimizer root:
-    nox -s "unit-3.12(torch_211-tf_latest)"
+    nox -s "unit-3.12(torch_211, tf_latest)"
 """
 
 import os

From 03dad8f90249c9141e8a2beae53fc0821ac2f4a8 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 03:40:12 -0700
Subject: [PATCH 5/9] Move regression tests to separate workflow file and
 folder name

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/bump_uv_lock.yml            |   3 +-
 .github/workflows/code_quality.yml            |   6 +-
 .github/workflows/gpu_tests.yml               |   8 --
 .github/workflows/pages.yml                   |  10 +-
 .github/workflows/regression_tests.yml        | 120 ++++++++++++++++++
 noxfile.py                                    |  12 +-
 .../torch/speculative/test_dflash.py          |   0
 7 files changed, 138 insertions(+), 21 deletions(-)
 create mode 100644 .github/workflows/regression_tests.yml
 rename tests/{gpu_regression => regression}/torch/speculative/test_dflash.py (100%)

diff --git a/.github/workflows/bump_uv_lock.yml b/.github/workflows/bump_uv_lock.yml
index 51f84e5c277..e0933418b6e 100644
--- a/.github/workflows/bump_uv_lock.yml
+++ b/.github/workflows/bump_uv_lock.yml
@@ -3,7 +3,8 @@ name: Bump uv.lock
 on:
   schedule:
     - cron: "0 9 * * 1" # Every Monday at 9:00 UTC
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
 
 permissions:
   contents: write
diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml
index cd877d18d24..3330d87332a 100644
--- a/.github/workflows/code_quality.yml
+++ b/.github/workflows/code_quality.yml
@@ -5,10 +5,12 @@ on:
     branches: [main, release/*, feature/*]
   schedule:
     - cron: "0 0 * * *" # Nightly
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
+
 
-# Cancel previous runs if new commit is pushed to the same PR
 concurrency:
+  # Cancel previous runs if new commit is pushed to the same PR
   group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
   cancel-in-progress: true
 
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index ed4fbe9de07..e4a8d3a1b07 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -48,11 +48,6 @@ jobs:
             tests/gpu/**
             tests/gpu_megatron/**
             tests/gpu_trtllm/**
-            tests/gpu_regression/**
-            examples/speculative_decoding/**
-            examples/dataset/**
-            modelopt_recipes/general/speculative_decoding/**
-            tools/launcher/**
           fail_on_initial_diff_error: true
   wait-checks:
     needs: [check-file-changes]
@@ -75,9 +70,6 @@ jobs:
             timeout: 60
             container_image: pytorch:26.01-py3
             # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers until https://github.com/tbenthompson/cppimport/pull/98
-          - example: gpu_regression
-            timeout: 15
-            container_image: pytorch:26.01-py3
           - example: gpu_megatron
             timeout: 45
             container_image: nemo:26.04
diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml
index eda19778594..43e2b5cc78d 100644
--- a/.github/workflows/pages.yml
+++ b/.github/workflows/pages.yml
@@ -8,16 +8,18 @@ on:
     branches: [main]
   schedule:
     - cron: "0 0 * * *" # Nightly
-  workflow_dispatch: # On-demand
+  workflow_dispatch:
+    # On-demand
+
 
-# Cancel previous runs if new commit is pushed
 concurrency:
+  # Cancel previous runs if new commit is pushed
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
   cancel-in-progress: true
 
 permissions:
-  contents: write       # push to gh-pages branch
-  pull-requests: write  # post/update preview URL comment on PRs
+  contents: write # push to gh-pages branch
+  pull-requests: write # post/update preview URL comment on PRs
 
 jobs:
   build-docs:
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
new file mode 100644
index 00000000000..e695e089a15
--- /dev/null
+++ b/.github/workflows/regression_tests.yml
@@ -0,0 +1,120 @@
+name: Regression tests
+
+on:
+  push:
+    branches: ["pull-request/[0-9]+"]
+    # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used
+  schedule:
+    - cron: "0 0 * * *" # Nightly
+  workflow_dispatch:
+    # On-demand
+
+
+concurrency:
+  # Cancel previous runs if new commit is pushed to the same PR
+  group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  check-file-changes:
+    if: startsWith(github.ref, 'refs/heads/pull-request/')
+    runs-on: ubuntu-latest
+    outputs:
+      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - id: get-pr-info
+        uses: nv-gha-runners/get-pr-info@main
+      # Get commit from main branch that is present in the PR to use as base for changed files
+      - id: calculate-merge-base
+        env:
+          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
+          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
+        run: |
+          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
+      - name: Check for changes in test-relevant directories
+        id: changed-tests
+        uses: step-security/changed-files@v46.0.5
+        with:
+          base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
+          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
+          files: |
+            .github/workflows/regression_tests.yml
+            modelopt/torch/**
+            noxfile.py
+            pyproject.toml
+            tests/regression/**
+            examples/speculative_decoding/**
+            examples/dataset/**
+            modelopt_recipes/general/speculative_decoding/**
+            tools/launcher/**
+          fail_on_initial_diff_error: true
+  wait-checks:
+    needs: [check-file-changes]
+    if: needs.check-file-changes.outputs.any_changed == 'true'
+    uses: ./.github/workflows/_wait_for_checks.yml
+    permissions:
+      checks: read
+    secrets: inherit
+    with:
+      match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass
+      delay: 300s
+  regression-tests-pr:
+    needs: [check-file-changes, wait-checks]
+    if: needs.check-file-changes.outputs.any_changed == 'true'
+    strategy: &regression_strategy
+      fail-fast: false
+      matrix:
+        include:
+          - example: regression
+            timeout: 15
+            container_image: pytorch:26.01-py3
+    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
+    timeout-minutes: ${{ matrix.timeout }}
+    container: &regression_container
+      image: nvcr.io/nvidia/${{ matrix.container_image }}
+      env:
+        GIT_DEPTH: 1000 # For correct version
+        PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+    steps: &regression_steps
+      - uses: actions/checkout@v6
+      - uses: nv-gha-runners/setup-proxy-cache@main
+      - name: Setup environment variables
+        run: |
+          echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
+      - name: Run regression tests
+        env:
+          COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
+          COVERAGE_FILE: ${{ github.workspace }}/.coverage
+        run: |
+          pip install nox
+          nox -s ${{ matrix.example }}
+      - name: Upload regression coverage to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: coverage.xml
+          flags: regression
+          fail_ci_if_error: false # test may be skipped if relevant file changes are not detected
+          verbose: true
+  regression-tests-non-pr:
+    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
+    strategy: *regression_strategy
+    runs-on: linux-amd64-gpu-rtxpro6000-latest-2
+    timeout-minutes: ${{ matrix.timeout }}
+    container: *regression_container
+    steps: *regression_steps
+  regression-pr-required-check:
+    # Run even if regression-tests-pr is skipped
+    if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
+    needs: [check-file-changes, regression-tests-pr]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Required regression tests did not succeed
+        if: |
+          needs.check-file-changes.result != 'success' ||
+          (needs.check-file-changes.outputs.any_changed == 'true' && needs.regression-tests-pr.result != 'success')
+        run: exit 1
diff --git a/noxfile.py b/noxfile.py
index 5cd8f8b5ed3..ca76ec107a0 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -128,18 +128,18 @@ def gpu_megatron(session):
     session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args())
 
 
-# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
+# Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 or later
 @nox.session(venv_backend="none")
-def gpu_regression(session):
+def gpu_trtllm(session):
     session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
-    session.run("python", "-m", "pytest", "tests/gpu_regression", *_cov_args())
+    session.run("python", "-m", "pytest", "tests/gpu_trtllm", *_cov_args())
 
 
-# Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 or later
+# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
 @nox.session(venv_backend="none")
-def gpu_trtllm(session):
+def regression(session):
     session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
-    session.run("python", "-m", "pytest", "tests/gpu_trtllm", *_cov_args())
+    session.run("python", "-m", "pytest", "tests/regression", *_cov_args())
 
 
 # ─── Code quality ─────────────────────────────────────────────────────────────
diff --git a/tests/gpu_regression/torch/speculative/test_dflash.py b/tests/regression/torch/speculative/test_dflash.py
similarity index 100%
rename from tests/gpu_regression/torch/speculative/test_dflash.py
rename to tests/regression/torch/speculative/test_dflash.py

From 16c827a5bb799291973d5265c58ce399a5fea1e0 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 03:44:27 -0700
Subject: [PATCH 6/9] update uv.lock

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 uv.lock | 426 +++++++++++++++-----------------------------------------
 1 file changed, 113 insertions(+), 313 deletions(-)

diff --git a/uv.lock b/uv.lock
index cfa742f1081..e223ce17eae 100644
--- a/uv.lock
+++ b/uv.lock
@@ -32,6 +32,9 @@ resolution-markers = [
     "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
 ]
 
+[manifest]
+overrides = [{ name = "torch", marker = "sys_platform == 'never'" }]
+
 [[package]]
 name = "accelerate"
 version = "1.13.0"
@@ -44,7 +47,7 @@ dependencies = [
     { name = "psutil" },
     { name = "pyyaml" },
     { name = "safetensors" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" }
 wheels = [
@@ -206,6 +209,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" },
 ]
 
+[[package]]
+name = "argcomplete"
+version = "3.6.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/38/61/0b9ae6399dd4a58d8c1b1dc5a27d6f2808023d0b5dd3104bb99f45a33ff6/argcomplete-3.6.3.tar.gz", hash = "sha256:62e8ed4fd6a45864acc8235409461b72c9a28ee785a2011cc5eb78318786c89c", size = 73754, upload-time = "2025-10-20T03:33:34.741Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/74/f5/9373290775639cb67a2fce7f629a1c240dce9f12fe927bc32b2736e16dfc/argcomplete-3.6.3-py3-none-any.whl", hash = "sha256:f5007b3a600ccac5d25bbce33089211dfd49eab4a7718da3f10e3082525a92ce", size = 43846, upload-time = "2025-10-20T03:33:33.021Z" },
+]
+
 [[package]]
 name = "async-timeout"
 version = "5.0.1"
@@ -266,15 +278,6 @@ toml = [
     { name = "tomli", marker = "python_full_version < '3.11'" },
 ]
 
-[[package]]
-name = "cachetools"
-version = "7.0.5"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/af/dd/57fe3fdb6e65b25a5987fd2cdc7e22db0aef508b91634d2e57d22928d41b/cachetools-7.0.5.tar.gz", hash = "sha256:0cd042c24377200c1dcd225f8b7b12b0ca53cc2c961b43757e774ebe190fd990", size = 37367, upload-time = "2026-03-09T20:51:29.451Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/06/f3/39cf3367b8107baa44f861dc802cbf16263c945b62d8265d36034fc07bea/cachetools-7.0.5-py3-none-any.whl", hash = "sha256:46bc8ebefbe485407621d0a4264b23c080cedd913921bad7ac3ed2f26c183114", size = 13918, upload-time = "2026-03-09T20:51:27.33Z" },
-]
-
 [[package]]
 name = "cbcbox"
 version = "2.929"
@@ -471,6 +474,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934", size = 46018, upload-time = "2021-06-11T10:22:42.561Z" },
 ]
 
+[[package]]
+name = "colorlog"
+version = "6.10.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/6d/c1/e419ef3723a074172b68aaa89c9f3de486ed4c2399e2dbd8113a4fdcaf9e/colorlog-6.10.1-py3-none-any.whl", hash = "sha256:2d7e8348291948af66122cff006c9f8da6255d224e7cf8e37d8de2df3bad8c9c", size = 11743, upload-time = "2025-10-16T16:14:10.512Z" },
+]
+
 [[package]]
 name = "coverage"
 version = "7.13.5"
@@ -570,21 +585,6 @@ dependencies = [
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/54/27/01d9078a77b9e31b79b9716e66ca4db74f4744c5232bcb3e8769395c4280/cppimport-22.8.2.tar.gz", hash = "sha256:bbb4957102db41bc99ad72c233bce92f9d1fd91be352fc07878c4361033a401f", size = 26635, upload-time = "2022-08-02T16:50:36.872Z" }
 
-[[package]]
-name = "cuda-bindings"
-version = "12.9.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cuda-pathfinder", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7a/d8/b546104b8da3f562c1ff8ab36d130c8fe1dd6a045ced80b4f6ad74f7d4e1/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d3c842c2a4303b2a580fe955018e31aea30278be19795ae05226235268032e5", size = 12148218, upload-time = "2025-10-21T14:51:28.855Z" },
-    { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" },
-    { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" },
-    { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" },
-    { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" },
-]
-
 [[package]]
 name = "cuda-pathfinder"
 version = "1.5.3"
@@ -659,11 +659,24 @@ dependencies = [
     { name = "psutil", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
     { name = "py-cpuinfo", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
     { name = "pydantic", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
-    { name = "torch", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "torch", marker = "sys_platform == 'never'" },
     { name = "tqdm", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/24/61/5ea1c63b139fe7530b196b68ce0bdffa9cde79882e527dcecae58bd6c770/deepspeed-0.18.9.tar.gz", hash = "sha256:ee4818dcf342794f74f429a0aeebef90291ec808fa82609c5140c23e665c4011", size = 1663466, upload-time = "2026-03-30T16:43:16.566Z" }
 
+[[package]]
+name = "dependency-groups"
+version = "1.3.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "packaging" },
+    { name = "tomli", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/62/55/f054de99871e7beb81935dea8a10b90cd5ce42122b1c3081d5282fdb3621/dependency_groups-1.3.1.tar.gz", hash = "sha256:78078301090517fd938c19f64a53ce98c32834dfe0dee6b88004a569a6adfefd", size = 10093, upload-time = "2025-05-02T00:34:29.452Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/99/c7/d1ec24fb280caa5a79b6b950db565dab30210a66259d17d5bb2b3a9f878d/dependency_groups-1.3.1-py3-none-any.whl", hash = "sha256:51aeaa0dfad72430fcfb7bcdbefbd75f3792e5919563077f30bc0d73f4493030", size = 8664, upload-time = "2025-05-02T00:34:27.085Z" },
+]
+
 [[package]]
 name = "diffusers"
 version = "0.37.1"
@@ -967,6 +980,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794, upload-time = "2021-09-17T21:40:39.897Z" },
 ]
 
+[[package]]
+name = "humanize"
+version = "4.15.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ba/66/a3921783d54be8a6870ac4ccffcd15c4dc0dd7fcce51c6d63b8c63935276/humanize-4.15.0.tar.gz", hash = "sha256:1dd098483eb1c7ee8e32eb2e99ad1910baefa4b75c3aff3a82f4d78688993b10", size = 83599, upload-time = "2025-12-20T20:16:13.19Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" },
+]
+
 [[package]]
 name = "hydra-core"
 version = "1.3.2"
@@ -1532,11 +1554,7 @@ name = "networkx"
 version = "3.4.2"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
-    "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'win32'",
     "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')",
-    "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'darwin'",
-    "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'",
-    "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'win32'",
     "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
 ]
@@ -1550,27 +1568,15 @@ name = "networkx"
 version = "3.6.1"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
-    "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'win32'",
-    "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'win32'",
-    "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'win32'",
     "(python_full_version >= '3.13' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')",
-    "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'darwin'",
     "(python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')",
-    "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'darwin'",
     "(python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')",
-    "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'darwin'",
     "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
     "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'",
-    "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'",
-    "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'win32'",
-    "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'",
-    "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'",
-    "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'",
-    "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'",
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" }
 wheels = [
@@ -1627,6 +1633,25 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" },
 ]
 
+[[package]]
+name = "nox"
+version = "2026.4.10"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "argcomplete" },
+    { name = "attrs" },
+    { name = "colorlog" },
+    { name = "dependency-groups" },
+    { name = "humanize" },
+    { name = "packaging" },
+    { name = "tomli", marker = "python_full_version < '3.11'" },
+    { name = "virtualenv" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e7/6b/e672c862a43cfca704d32359221fa3780226daa1e5db5dfc401bcc8be9c9/nox-2026.4.10.tar.gz", hash = "sha256:2d0af5374f3f37a295428c927d1b04a8182aa01762897d172446dda2f1ce9692", size = 4034839, upload-time = "2026-04-10T17:42:42.209Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7f/95/4df134a100b5a9a12378d5301b934366686ef6fbdaffcd21211d5654970e/nox-2026.4.10-py3-none-any.whl", hash = "sha256:082c117627590d9b90aa21f86df89b310b07c5842539524203bcb3c719f116c1", size = 75536, upload-time = "2026-04-10T17:42:40.664Z" },
+]
+
 [[package]]
 name = "numpy"
 version = "2.2.6"
@@ -1779,108 +1804,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/04/74/f4c001f4714c3ad9ce037e18cf2b9c64871a84951eaa0baf683a9ca9301c/numpy-2.4.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:f2cf083b324a467e1ab358c105f6cad5ea950f50524668a80c486ff1db24e119", size = 12509075, upload-time = "2026-03-29T13:21:57.644Z" },
 ]
 
-[[package]]
-name = "nvidia-cublas-cu12"
-version = "12.8.4.1"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" },
-]
-
-[[package]]
-name = "nvidia-cuda-cupti-cu12"
-version = "12.8.90"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" },
-]
-
-[[package]]
-name = "nvidia-cuda-nvrtc-cu12"
-version = "12.8.93"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" },
-]
-
-[[package]]
-name = "nvidia-cuda-runtime-cu12"
-version = "12.8.90"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" },
-]
-
-[[package]]
-name = "nvidia-cudnn-cu12"
-version = "9.10.2.21"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" },
-]
-
-[[package]]
-name = "nvidia-cufft-cu12"
-version = "11.3.3.83"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" },
-]
-
-[[package]]
-name = "nvidia-cufile-cu12"
-version = "1.13.1.3"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" },
-]
-
-[[package]]
-name = "nvidia-curand-cu12"
-version = "10.3.9.90"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" },
-]
-
-[[package]]
-name = "nvidia-cusolver-cu12"
-version = "11.7.3.90"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-    { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" },
-]
-
-[[package]]
-name = "nvidia-cusparse-cu12"
-version = "12.5.8.93"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" },
-]
-
-[[package]]
-name = "nvidia-cusparselt-cu12"
-version = "0.7.1"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" },
-]
-
 [[package]]
 name = "nvidia-ml-py"
 version = "13.595.45"
@@ -1909,7 +1832,7 @@ dependencies = [
     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "setuptools" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
     { name = "tqdm" },
 ]
 
@@ -1969,6 +1892,7 @@ dev = [
     { name = "ml-dtypes" },
     { name = "mypy" },
     { name = "nltk" },
+    { name = "nox" },
     { name = "onnx" },
     { name = "onnx-graphsurgeon" },
     { name = "onnxconverter-common" },
@@ -2002,10 +1926,9 @@ dev = [
     { name = "torch-geometric" },
     { name = "torchprofile" },
     { name = "torchvision" },
-    { name = "tox" },
-    { name = "tox-current-env" },
     { name = "transformers" },
     { name = "typeguard" },
+    { name = "uv" },
     { name = "wonderwords" },
 ]
 dev-docs = [
@@ -2027,6 +1950,7 @@ dev-lint = [
 ]
 dev-test = [
     { name = "coverage", extra = ["toml"] },
+    { name = "nox" },
     { name = "pytest" },
     { name = "pytest-cov" },
     { name = "pytest-instafail" },
@@ -2035,8 +1959,7 @@ dev-test = [
     { name = "torch-geometric" },
     { name = "torchprofile" },
     { name = "torchvision" },
-    { name = "tox" },
-    { name = "tox-current-env" },
+    { name = "uv" },
 ]
 hf = [
     { name = "accelerate" },
@@ -2100,6 +2023,7 @@ requires-dist = [
     { name = "mypy", marker = "extra == 'dev-lint'", specifier = "==1.17.1" },
     { name = "ninja" },
     { name = "nltk", marker = "extra == 'hf'" },
+    { name = "nox", marker = "extra == 'dev-test'" },
     { name = "numpy" },
     { name = "nvidia-ml-py", specifier = ">=12" },
     { name = "nvidia-modelopt", extras = ["all", "dev-docs", "dev-lint", "dev-test"], marker = "extra == 'dev'" },
@@ -2120,7 +2044,7 @@ requires-dist = [
     { name = "peft", marker = "extra == 'hf'", specifier = ">=0.17.0" },
     { name = "polygraphy", marker = "extra == 'onnx'", specifier = ">=0.49.22" },
     { name = "pre-commit", marker = "extra == 'dev-lint'", specifier = "==4.3.0" },
-    { name = "pulp" },
+    { name = "pulp", specifier = "<4.0" },
     { name = "pydantic", specifier = ">=2.0" },
     { name = "pytest", marker = "extra == 'dev-test'" },
     { name = "pytest-cov", marker = "extra == 'dev-test'" },
@@ -2147,47 +2071,14 @@ requires-dist = [
     { name = "torch-geometric", marker = "extra == 'dev-test'" },
     { name = "torchprofile", marker = "extra == 'dev-test'", specifier = "==0.0.4" },
     { name = "torchvision", marker = "extra == 'dev-test'" },
-    { name = "tox", marker = "extra == 'dev-test'", specifier = ">4.18" },
-    { name = "tox-current-env", marker = "extra == 'dev-test'", specifier = ">=0.0.12" },
     { name = "tqdm" },
     { name = "transformers", marker = "extra == 'hf'", specifier = ">=4.56" },
     { name = "typeguard", marker = "extra == 'puzzletron'" },
+    { name = "uv", marker = "extra == 'dev-test'" },
     { name = "wonderwords", marker = "extra == 'hf'" },
 ]
 provides-extras = ["onnx", "hf", "puzzletron", "dev-lint", "dev-docs", "dev-test", "all", "dev"]
 
-[[package]]
-name = "nvidia-nccl-cu12"
-version = "2.27.5"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" },
-]
-
-[[package]]
-name = "nvidia-nvjitlink-cu12"
-version = "12.8.93"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" },
-]
-
-[[package]]
-name = "nvidia-nvshmem-cu12"
-version = "3.4.5"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" },
-]
-
-[[package]]
-name = "nvidia-nvtx-cu12"
-version = "12.8.90"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" },
-]
-
 [[package]]
 name = "omegaconf"
 version = "2.3.0"
@@ -2604,7 +2495,7 @@ dependencies = [
     { name = "psutil" },
     { name = "pyyaml" },
     { name = "safetensors" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
     { name = "tqdm" },
     { name = "transformers" },
 ]
@@ -3061,19 +2952,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" },
 ]
 
-[[package]]
-name = "pyproject-api"
-version = "1.10.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "packaging" },
-    { name = "tomli", marker = "python_full_version < '3.11'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/45/7b/c0e1333b61d41c69e59e5366e727b18c4992688caf0de1be10b3e5265f6b/pyproject_api-1.10.0.tar.gz", hash = "sha256:40c6f2d82eebdc4afee61c773ed208c04c19db4c4a60d97f8d7be3ebc0bbb330", size = 22785, upload-time = "2025-10-09T19:12:27.21Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/54/cc/cecf97be298bee2b2a37dd360618c819a2a7fd95251d8e480c1f0eb88f3b/pyproject_api-1.10.0-py3-none-any.whl", hash = "sha256:8757c41a79c0f4ab71b99abed52b97ecf66bd20b04fa59da43b5840bac105a09", size = 13218, upload-time = "2025-10-09T19:12:24.428Z" },
-]
-
 [[package]]
 name = "pyreadline3"
 version = "3.5.4"
@@ -3944,7 +3822,7 @@ dependencies = [
     { name = "huggingface-hub" },
     { name = "pyyaml" },
     { name = "safetensors" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
     { name = "torchvision" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/7b/1e/e924b3b2326a856aaf68586f9c52a5fc81ef45715eca408393b68c597e0e/timm-1.0.26.tar.gz", hash = "sha256:f66f082f2f381cf68431c22714c8b70f723837fa2a185b155961eab90f2d5b10", size = 2419859, upload-time = "2026-03-23T18:12:10.272Z" }
@@ -4018,76 +3896,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" },
 ]
 
-[[package]]
-name = "tomli-w"
-version = "1.2.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/19/75/241269d1da26b624c0d5e110e8149093c759b7a286138f4efd61a60e75fe/tomli_w-1.2.0.tar.gz", hash = "sha256:2dd14fac5a47c27be9cd4c976af5a12d87fb1f0b4512f81d69cce3b35ae25021", size = 7184, upload-time = "2025-01-15T12:07:24.262Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c7/18/c86eb8e0202e32dd3df50d43d7ff9854f8e0603945ff398974c1d91ac1ef/tomli_w-1.2.0-py3-none-any.whl", hash = "sha256:188306098d013b691fcadc011abd66727d3c414c571bb01b1a174ba8c983cf90", size = 6675, upload-time = "2025-01-15T12:07:22.074Z" },
-]
-
 [[package]]
 name = "torch"
 version = "2.10.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "filelock" },
-    { name = "fsspec" },
-    { name = "jinja2" },
-    { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
-    { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "setuptools", marker = "python_full_version >= '3.12'" },
-    { name = "sympy" },
-    { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "typing-extensions" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5b/30/bfebdd8ec77db9a79775121789992d6b3b75ee5494971294d7b4b7c999bc/torch-2.10.0-2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:2b980edd8d7c0a68c4e951ee1856334a43193f98730d97408fbd148c1a933313", size = 79411457, upload-time = "2026-02-10T21:44:59.189Z" },
-    { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" },
-    { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" },
-    { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" },
-    { url = "https://files.pythonhosted.org/packages/16/ee/efbd56687be60ef9af0c9c0ebe106964c07400eade5b0af8902a1d8cd58c/torch-2.10.0-3-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a1ff626b884f8c4e897c4c33782bdacdff842a165fee79817b1dd549fdda1321", size = 915510070, upload-time = "2026-03-11T14:16:39.386Z" },
-    { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" },
-    { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" },
-    { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" },
-    { url = "https://files.pythonhosted.org/packages/0c/1a/c61f36cfd446170ec27b3a4984f072fd06dab6b5d7ce27e11adb35d6c838/torch-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5276fa790a666ee8becaffff8acb711922252521b28fbce5db7db5cf9cb2026d", size = 145992962, upload-time = "2026-01-21T16:24:14.04Z" },
-    { url = "https://files.pythonhosted.org/packages/b5/60/6662535354191e2d1555296045b63e4279e5a9dbad49acf55a5d38655a39/torch-2.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:aaf663927bcd490ae971469a624c322202a2a1e68936eb952535ca4cd3b90444", size = 915599237, upload-time = "2026-01-21T16:23:25.497Z" },
-    { url = "https://files.pythonhosted.org/packages/40/b8/66bbe96f0d79be2b5c697b2e0b187ed792a15c6c4b8904613454651db848/torch-2.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:a4be6a2a190b32ff5c8002a0977a25ea60e64f7ba46b1be37093c141d9c49aeb", size = 113720931, upload-time = "2026-01-21T16:24:23.743Z" },
-    { url = "https://files.pythonhosted.org/packages/76/bb/d820f90e69cda6c8169b32a0c6a3ab7b17bf7990b8f2c680077c24a3c14c/torch-2.10.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:35e407430795c8d3edb07a1d711c41cc1f9eaddc8b2f1cc0a165a6767a8fb73d", size = 79411450, upload-time = "2026-01-21T16:25:30.692Z" },
-    { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" },
-    { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" },
-    { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" },
-    { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" },
-    { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" },
-    { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" },
-    { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" },
-    { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" },
-    { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" },
-    { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" },
-    { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" },
-    { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" },
-    { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" },
-    { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" },
+    { name = "filelock", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "fsspec", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "jinja2", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "setuptools", marker = "(python_full_version >= '3.12' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.12' and sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "sympy", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
+    { name = "typing-extensions", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" },
 ]
 
 [[package]]
@@ -4118,7 +3939,7 @@ source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
     { name = "torchvision" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/6f/36/574c0c46e818533b78b3c09505211162918188325ab4165ef11a3f295755/torchprofile-0.0.4.tar.gz", hash = "sha256:96b6da17d752a06b02977e078aea95614893b31d4117dd5dcd081f30ce65611b", size = 4557, upload-time = "2021-06-22T04:58:03.592Z" }
@@ -4134,7 +3955,7 @@ dependencies = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "pillow" },
-    { name = "torch" },
+    { name = "torch", marker = "sys_platform == 'never'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/50/ae/cbf727421eb73f1cf907fbe5788326a08f111b3f6b6ddca15426b53fec9a/torchvision-0.25.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a95c47abb817d4e90ea1a8e57bd0d728e3e6b533b3495ae77d84d883c4d11f56", size = 1874919, upload-time = "2026-01-21T16:27:47.617Z" },
@@ -4159,41 +3980,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" },
 ]
 
-[[package]]
-name = "tox"
-version = "4.53.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cachetools" },
-    { name = "colorama" },
-    { name = "filelock" },
-    { name = "packaging" },
-    { name = "platformdirs" },
-    { name = "pluggy" },
-    { name = "pyproject-api" },
-    { name = "python-discovery" },
-    { name = "tomli", marker = "python_full_version < '3.11'" },
-    { name = "tomli-w" },
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
-    { name = "virtualenv" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/04/01/d87a00063fa670ce4c48a9706b615a95ddf2c9ef5558d43af6071f166fd4/tox-4.53.0.tar.gz", hash = "sha256:62c780e42f87d34ee60f2ea20342156253794fdcbd6885fd797d98ee05009f22", size = 274048, upload-time = "2026-04-14T13:44:13.782Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/16/03/02e2a03f3756cfb66e7e1bac41b06953f12cec75ddb961d56695d4d43dc4/tox-4.53.0-py3-none-any.whl", hash = "sha256:cc4e716d18c4889aa179d785175c438fa60c35deef20ce689ec288d8fb656096", size = 212164, upload-time = "2026-04-14T13:44:11.997Z" },
-]
-
-[[package]]
-name = "tox-current-env"
-version = "0.0.16"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "tox" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/35/f3/504ab508410e0660d989eb3ea175f6245271c5582dc4a93268fa9feceeda/tox_current_env-0.0.16.tar.gz", hash = "sha256:2e453c3e82e837d35846004a678db4504e24e5c0419d6e42aa07ca8294fad1bd", size = 25129, upload-time = "2025-03-12T15:52:14.845Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/05/38/b485aed12bb714c53aedbed4919d7ffda7a0f21760229c894feedcc05a61/tox_current_env-0.0.16-py3-none-any.whl", hash = "sha256:88db23525a6514dee6350df01f5b591654b72e7709df3c57e9e1edb1bbc22737", size = 13543, upload-time = "2025-03-12T15:52:13.226Z" },
-]
-
 [[package]]
 name = "tqdm"
 version = "4.67.3"
@@ -4239,18 +4025,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" },
 ]
 
-[[package]]
-name = "triton"
-version = "3.6.0"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/8c/f7/f1c9d3424ab199ac53c2da567b859bcddbb9c9e7154805119f8bd95ec36f/triton-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a6550fae429e0667e397e5de64b332d1e5695b73650ee75a6146e2e902770bea", size = 188105201, upload-time = "2026-01-20T16:00:29.272Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" },
-    { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" },
-    { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" },
-]
-
 [[package]]
 name = "typer"
 version = "0.24.1"
@@ -4305,6 +4079,32 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" },
 ]
 
+[[package]]
+name = "uv"
+version = "0.11.7"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/9b/7d/17750123a8c8e324627534fe1ae2e7a46689db8492f1a834ab4fd229a7d8/uv-0.11.7.tar.gz", hash = "sha256:46d971489b00bdb27e0aa715e4a5cd4ef2c28ea5b6ef78f2b67bf861eb44b405", size = 4083385, upload-time = "2026-04-15T21:42:55.474Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b2/5b/2bb2ab6fe6c78c2be10852482ef0cae5f3171460a6e5e24c32c9a0843163/uv-0.11.7-py3-none-linux_armv6l.whl", hash = "sha256:f422d39530516b1dfb28bb6e90c32bb7dacd50f6a383cd6e40c1a859419fbc8c", size = 23757265, upload-time = "2026-04-15T21:43:14.494Z" },
+    { url = "https://files.pythonhosted.org/packages/b2/f5/36ff27b01e60a88712628c8a5a6003b8e418883c24e084e506095844a797/uv-0.11.7-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:8b2fe1ec6775dad10183e3fdce430a5b37b7857d49763c884f3a67eaa8ca6f8a", size = 23184529, upload-time = "2026-04-15T21:42:30.225Z" },
+    { url = "https://files.pythonhosted.org/packages/8a/fa/f379be661316698f877e78f4c51e5044be0b6f390803387237ad92c4057f/uv-0.11.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:162fa961a9a081dcea6e889c79f738a5ae56507047e4672964972e33c301bea9", size = 21780167, upload-time = "2026-04-15T21:42:44.942Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/7f/fbed29775b0612f4f5679d3226268f1a347161abc1727b4080fb41d9f46f/uv-0.11.7-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:5985a15a92bd9a170fc1947abb1fbc3e9828c5a430ad85b5bed8356c20b67a71", size = 23609640, upload-time = "2026-04-15T21:42:22.57Z" },
+    { url = "https://files.pythonhosted.org/packages/ad/de/989a69634a869a22322770120557c2d8cbba5b77ec7cfad326b4ec0f0547/uv-0.11.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:fab0bb43fbbc0ee5b5fee212078d2300c371b725faff7cf72eeaafa0bff0606b", size = 23322484, upload-time = "2026-04-15T21:43:26.52Z" },
+    { url = "https://files.pythonhosted.org/packages/24/08/c1af05ea602eb4eb75d86badb6b0594cc104c3ca83ccf06d9ed4dd2186ad/uv-0.11.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:23d457d6731ebdb83f1bffebe4894edab2ef43c1ec5488433c74300db4958924", size = 23326385, upload-time = "2026-04-15T21:42:41.32Z" },
+    { url = "https://files.pythonhosted.org/packages/68/99/e246962da06383e992ecab55000c62a50fb36efef855ea7264fad4816bf4/uv-0.11.7-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7d6a17507b8139b8803f445a03fd097f732ce8356b1b7b13cdb4dd8ef7f4b2e0", size = 24985751, upload-time = "2026-04-15T21:42:37.777Z" },
+    { url = "https://files.pythonhosted.org/packages/45/2d/b0b68083859579ce811996c1480765ec6a2442b44c451eaef53e6218fbae/uv-0.11.7-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dd48823ca4b505124389f49ae50626ba9f57212b9047738efc95126ed5f3844d", size = 25724160, upload-time = "2026-04-15T21:43:18.762Z" },
+    { url = "https://files.pythonhosted.org/packages/4e/19/5970e89d9e458fd3c4966bbc586a685a1c0ab0a8bf334503f63fa20b925b/uv-0.11.7-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:eb91f52ee67e10d5290f2c2897e2171357f1a10966de38d83eefa93d96843b0c", size = 25028512, upload-time = "2026-04-15T21:43:02.721Z" },
+    { url = "https://files.pythonhosted.org/packages/83/eb/4e1557daf6693cb446ed28185664ad6682fd98c6dbac9e433cbc35df450a/uv-0.11.7-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e4d5e31bea86e1b6e0f5a0f95e14e80018e6f6c0129256d2915a4b3d793644d", size = 24933975, upload-time = "2026-04-15T21:42:18.828Z" },
+    { url = "https://files.pythonhosted.org/packages/68/55/3b517ec8297f110d6981f525cccf26f86e30883fbb9c282769cffbcdcfca/uv-0.11.7-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ceae53b202ea92bc954759bc7c7570cdcd5c3512fce15701198c19fd2dfb8605", size = 23706403, upload-time = "2026-04-15T21:43:10.664Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/30/7d93a0312d60e147722967036dc8ea37baab4802784bddc22464cb707deb/uv-0.11.7-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:f97e9f4e4d44fb5c4dfaa05e858ef3414a96416a2e4af270ecd88a3e5fb049a9", size = 24495797, upload-time = "2026-04-15T21:42:26.538Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/89/d49480bdab7725d36982793857e461d471bde8e1b7f438ffccee677a7bf8/uv-0.11.7-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:750ee5b96959b807cf442b73dd8b55111862d63f258f896787ea5f06b68aaca9", size = 24580471, upload-time = "2026-04-15T21:42:52.871Z" },
+    { url = "https://files.pythonhosted.org/packages/b6/9f/c57dc03b48be17b564e304eb9ff982890c12dfb888b1ce370788733329ab/uv-0.11.7-py3-none-musllinux_1_1_i686.whl", hash = "sha256:f394331f0507e80ee732cb3df737589de53bed999dd02a6d24682f08c2f8ac4f", size = 24113637, upload-time = "2026-04-15T21:42:34.094Z" },
+    { url = "https://files.pythonhosted.org/packages/13/ba/b87e358b629a68258527e3490e73b7b148770f4d2257842dea3b7981d4e8/uv-0.11.7-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:0df59ab0c6a4b14a763e8445e1c303af9abeb53cdfa4428daf9ff9642c0a3cce", size = 25119850, upload-time = "2026-04-15T21:43:22.529Z" },
+    { url = "https://files.pythonhosted.org/packages/4b/74/16d229e1d8574bcbafa6dc643ac20b70c3e581f42ac31a6f4fd53035ffe3/uv-0.11.7-py3-none-win32.whl", hash = "sha256:553e67cc766d013ce24353fecd4ea5533d2aedcfd35f9fac430e07b1d1f23ed4", size = 22918454, upload-time = "2026-04-15T21:42:58.702Z" },
+    { url = "https://files.pythonhosted.org/packages/a6/1d/b73e473da616ac758b8918fb218febcc46ddf64cba9e03894dfa226b28bd/uv-0.11.7-py3-none-win_amd64.whl", hash = "sha256:5674dfb5944513f4b3735b05c2deba6b1b01151f46729d533d413a9a905f8c5d", size = 25447744, upload-time = "2026-04-15T21:42:48.813Z" },
+    { url = "https://files.pythonhosted.org/packages/1b/bb/e6bfdea92ed270f3445a5a3c17599d041b3f2dbc5026c09e02830a03bbaf/uv-0.11.7-py3-none-win_arm64.whl", hash = "sha256:6158b7e39464f1aa1e040daa0186cae4749a78b5cd80ac769f32ca711b8976b1", size = 23941816, upload-time = "2026-04-15T21:43:06.732Z" },
+]
+
 [[package]]
 name = "uvicorn"
 version = "0.44.0"

From 7ac16693fa65afd5e9eacd32ffa80a4d138a2891 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 06:25:32 -0700
Subject: [PATCH 7/9] fix ci

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 noxfile.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/noxfile.py b/noxfile.py
index ca76ec107a0..fcef3d30875 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -122,8 +122,8 @@ def gpu(session):
 # Container: nvcr.io/nvidia/nemo:26.04 or later
 @nox.session(venv_backend="none")
 def gpu_megatron(session):
-    # nemo:26.04 has transformers 5.x but tensorrt_llm 1.2.0 which does not support it causing import errors
-    session.run("python", "-m", "pip", "uninstall", "-y", "tensorrt_llm")
+    # nemo:26.04 has transformers 5.x but system-wide installed trtllm 1.2.0 which does not support it causing import errors
+    session.run("pip", "uninstall", "-y", "tensorrt_llm")
     session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]")
     session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args())
 

From 922dd02217c5df24281f9708caadf1ff0010327b Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 07:45:06 -0700
Subject: [PATCH 8/9] Simplify CI configuration

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/_pr_gate.yml         |  51 +++++++++++++
 .github/workflows/example_tests.yml    |  67 +++++------------
 .github/workflows/gpu_tests.yml        |  57 ++++----------
 .github/workflows/regression_tests.yml |  63 +++++-----------
 .github/workflows/unit_tests.yml       | 100 ++++++++++++-------------
 5 files changed, 150 insertions(+), 188 deletions(-)
 create mode 100644 .github/workflows/_pr_gate.yml

diff --git a/.github/workflows/_pr_gate.yml b/.github/workflows/_pr_gate.yml
new file mode 100644
index 00000000000..28afff3f64d
--- /dev/null
+++ b/.github/workflows/_pr_gate.yml
@@ -0,0 +1,51 @@
+name: PR Gate
+
+on:
+  workflow_call:
+    inputs:
+      files:
+        description: "Newline-separated list of file patterns to watch for changes"
+        required: true
+        type: string
+    outputs:
+      any_changed:
+        description: "Whether any relevant files changed"
+        value: ${{ jobs.check-file-changes.outputs.any_changed }}
+
+jobs:
+  check-file-changes:
+    if: startsWith(github.ref, 'refs/heads/pull-request/')
+    runs-on: ubuntu-latest
+    outputs:
+      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - id: get-pr-info
+        uses: nv-gha-runners/get-pr-info@main
+      # Get commit from main branch that is present in the PR to use as base for changed files
+      - id: calculate-merge-base
+        env:
+          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
+          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
+        run: |
+          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
+      - name: Check for changes in test-relevant directories
+        id: changed-tests
+        uses: step-security/changed-files@v46.0.5
+        with:
+          base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
+          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
+          files: ${{ inputs.files }}
+          fail_on_initial_diff_error: true
+  wait-checks:
+    needs: [check-file-changes]
+    if: needs.check-file-changes.outputs.any_changed == 'true'
+    uses: ./.github/workflows/_wait_for_checks.yml
+    permissions:
+      checks: read
+    secrets: inherit
+    with:
+      match_pattern: "^linux$" # Wait for Unit tests / linux (DCO is a prerequisite of linux)
+      delay: 300s
diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index cbccfaa6149..55696ea1570 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -16,52 +16,23 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  check-file-changes:
-    if: startsWith(github.ref, 'refs/heads/pull-request/')
-    runs-on: ubuntu-latest
-    outputs:
-      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-      - id: get-pr-info
-        uses: nv-gha-runners/get-pr-info@main
-      # Get commit from main branch that is present in the PR to use as base for changed files
-      - id: calculate-merge-base
-        env:
-          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
-        run: |
-          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
-      - name: Check for changes in test-relevant directories
-        id: changed-tests
-        uses: step-security/changed-files@v46.0.5
-        with:
-          base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
-          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          files: |
-            .github/workflows/example_tests.yml
-            examples/**
-            modelopt/**
-            pyproject.toml
-            tests/examples/**
-          fail_on_initial_diff_error: true
-  wait-checks:
-    needs: [check-file-changes]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
-    uses: ./.github/workflows/_wait_for_checks.yml
+  pr-gate:
+    uses: ./.github/workflows/_pr_gate.yml
     permissions:
       checks: read
     secrets: inherit
     with:
-      match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass
-      delay: 300s
+      files: |
+        .github/workflows/example_tests.yml
+        examples/**
+        modelopt/**
+        pyproject.toml
+        tests/examples/**
 
   ##### PyTorch Example Tests (speculative_decoding requires 26.01 image) #####
   torch-pr:
-    needs: [check-file-changes, wait-checks]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
     strategy: &torch_strategy
       fail-fast: false
       matrix:
@@ -92,8 +63,8 @@ jobs:
 
   ##### TensorRT-LLM Example Tests #####
   trtllm-pr:
-    needs: [check-file-changes, wait-checks]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
     strategy:
       fail-fast: false
       matrix:
@@ -122,8 +93,8 @@ jobs:
 
   ##### Megatron Example Tests #####
   megatron-pr:
-    needs: [check-file-changes, wait-checks]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
     strategy: &nemo_strategy
       fail-fast: false
       matrix:
@@ -151,8 +122,8 @@ jobs:
 
   ##### ONNX/TensorRT Example Tests #####
   onnx-pr:
-    needs: [check-file-changes, wait-checks]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
     strategy: &onnx_strategy
       fail-fast: false
       matrix:
@@ -180,13 +151,13 @@ jobs:
   example-pr-required-check:
     # Run even if example tests are skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [check-file-changes, torch-pr, trtllm-pr, megatron-pr, onnx-pr]
+    needs: [pr-gate, torch-pr, trtllm-pr, megatron-pr, onnx-pr]
     runs-on: ubuntu-latest
     steps:
       - name: Required GPU tests did not succeed
         if: |
-          needs.check-file-changes.result != 'success' ||
-          (needs.check-file-changes.outputs.any_changed == 'true' && (
+          needs.pr-gate.result != 'success' ||
+          (needs.pr-gate.outputs.any_changed == 'true' && (
             needs.torch-pr.result != 'success' ||
             needs.trtllm-pr.result != 'success' ||
             needs.megatron-pr.result != 'success' ||
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index e4a8d3a1b07..30a872e91b3 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -16,52 +16,23 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  check-file-changes:
-    if: startsWith(github.ref, 'refs/heads/pull-request/')
-    runs-on: ubuntu-latest
-    outputs:
-      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-      - id: get-pr-info
-        uses: nv-gha-runners/get-pr-info@main
-      # Get commit from main branch that is present in the PR to use as base for changed files
-      - id: calculate-merge-base
-        env:
-          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
-        run: |
-          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
-      - name: Check for changes in test-relevant directories
-        id: changed-tests
-        uses: step-security/changed-files@v46.0.5
-        with:
-          base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
-          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          files: |
-            .github/workflows/gpu_tests.yml
-            modelopt/**
-            noxfile.py
-            pyproject.toml
-            tests/gpu/**
-            tests/gpu_megatron/**
-            tests/gpu_trtllm/**
-          fail_on_initial_diff_error: true
-  wait-checks:
-    needs: [check-file-changes]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
-    uses: ./.github/workflows/_wait_for_checks.yml
+  pr-gate:
+    uses: ./.github/workflows/_pr_gate.yml
     permissions:
       checks: read
     secrets: inherit
     with:
-      match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass
-      delay: 300s
+      files: |
+        .github/workflows/gpu_tests.yml
+        modelopt/**
+        noxfile.py
+        pyproject.toml
+        tests/gpu/**
+        tests/gpu_megatron/**
+        tests/gpu_trtllm/**
   gpu-tests-pr:
-    needs: [check-file-changes, wait-checks]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: needs.pr-gate.outputs.any_changed == 'true'
     strategy: &gpu_strategy
       fail-fast: false
       matrix:
@@ -114,9 +85,9 @@ jobs:
   gpu-pr-required-check:
     # Run even if gpu-tests-pr is skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [check-file-changes, gpu-tests-pr]
+    needs: [pr-gate, gpu-tests-pr]
     runs-on: ubuntu-latest
     steps:
       - name: Required GPU tests did not succeed
-        if: ${{ needs.check-file-changes.result != 'success' || (needs.check-file-changes.outputs.any_changed == 'true' && needs.gpu-tests-pr.result != 'success') }}
+        if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests-pr.result != 'success') }}
         run: exit 1
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
index e695e089a15..2635d751b7b 100644
--- a/.github/workflows/regression_tests.yml
+++ b/.github/workflows/regression_tests.yml
@@ -16,54 +16,25 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  check-file-changes:
-    if: startsWith(github.ref, 'refs/heads/pull-request/')
-    runs-on: ubuntu-latest
-    outputs:
-      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-      - id: get-pr-info
-        uses: nv-gha-runners/get-pr-info@main
-      # Get commit from main branch that is present in the PR to use as base for changed files
-      - id: calculate-merge-base
-        env:
-          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
-        run: |
-          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
-      - name: Check for changes in test-relevant directories
-        id: changed-tests
-        uses: step-security/changed-files@v46.0.5
-        with:
-          base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
-          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          files: |
-            .github/workflows/regression_tests.yml
-            modelopt/torch/**
-            noxfile.py
-            pyproject.toml
-            tests/regression/**
-            examples/speculative_decoding/**
-            examples/dataset/**
-            modelopt_recipes/general/speculative_decoding/**
-            tools/launcher/**
-          fail_on_initial_diff_error: true
-  wait-checks:
-    needs: [check-file-changes]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
-    uses: ./.github/workflows/_wait_for_checks.yml
+  pr-gate:
+    uses: ./.github/workflows/_pr_gate.yml
     permissions:
       checks: read
     secrets: inherit
     with:
-      match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass
-      delay: 300s
+      files: |
+        .github/workflows/regression_tests.yml
+        modelopt/torch/**
+        noxfile.py
+        pyproject.toml
+        tests/regression/**
+        examples/speculative_decoding/**
+        examples/dataset/**
+        modelopt_recipes/general/speculative_decoding/**
+        tools/launcher/**
   regression-tests-pr:
-    needs: [check-file-changes, wait-checks]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [pr-gate]
+    if: needs.pr-gate.outputs.any_changed == 'true'
     strategy: &regression_strategy
       fail-fast: false
       matrix:
@@ -110,11 +81,11 @@ jobs:
   regression-pr-required-check:
     # Run even if regression-tests-pr is skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [check-file-changes, regression-tests-pr]
+    needs: [pr-gate, regression-tests-pr]
     runs-on: ubuntu-latest
     steps:
       - name: Required regression tests did not succeed
         if: |
-          needs.check-file-changes.result != 'success' ||
-          (needs.check-file-changes.outputs.any_changed == 'true' && needs.regression-tests-pr.result != 'success')
+          needs.pr-gate.result != 'success' ||
+          (needs.pr-gate.outputs.any_changed == 'true' && needs.regression-tests-pr.result != 'success')
         run: exit 1
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index 8401cb9120e..889cdbee83f 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -32,6 +32,26 @@ jobs:
     secrets: inherit
     with:
       match_pattern: "^DCO$"
+  check-file-changes:
+    if: github.event_name == 'pull_request'
+    runs-on: ubuntu-latest
+    outputs:
+      any_changed: ${{ steps.changed.outputs.any_changed }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - name: Check for changes in test-relevant paths
+        id: changed
+        uses: step-security/changed-files@v46.0.5
+        with:
+          files: |
+            .github/workflows/unit_tests.yml
+            modelopt/**
+            noxfile.py
+            pyproject.toml
+            tests/unit/**
+            tools/launcher/**
   linux:
     needs: [check-dco]
     runs-on: ubuntu-latest
@@ -52,8 +72,8 @@ jobs:
           fail_ci_if_error: true
           verbose: true
   windows:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
+    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [linux, check-file-changes]
     runs-on: windows-latest
     timeout-minutes: 30
     steps:
@@ -63,53 +83,46 @@ jobs:
           python-version: "3.12"
       - name: Run unit tests (without coverage)
         run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)"
-  multi-py:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
+  multi-version:
+    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:
-        py: [10, 11, 13]
+        include:
+          - {nox_session: "unit-3.10(torch_211, tf_latest)", python_version: "3.10"}
+          - {nox_session: "unit-3.11(torch_211, tf_latest)", python_version: "3.11"}
+          - {nox_session: "unit-3.13(torch_211, tf_latest)", python_version: "3.13"}
+          - {nox_session: "unit-3.12(torch_28, tf_latest)", python_version: "3.12"}
+          - {nox_session: "unit-3.12(torch_29, tf_latest)", python_version: "3.12"}
+          - {nox_session: "unit-3.12(torch_210, tf_latest)", python_version: "3.12"}
+          - {nox_session: "unit-3.12(torch_211, tf_min)", python_version: "3.12"}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
         with:
-          python-version: "3.${{ matrix.py }}"
-      - name: Run unit tests
-        run: pip install nox uv && nox -s "unit-3.${{ matrix.py }}(torch_211, tf_latest)"
-  multi-torch:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
-    runs-on: ubuntu-latest
-    timeout-minutes: 30
-    strategy:
-      fail-fast: false
-      matrix:
-        torch: [28, 29, 210]
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/ubuntu-setup
+          python-version: ${{ matrix.python_version }}
       - name: Run unit tests
-        run: pip install nox uv && nox -s "unit-3.12(torch_${{ matrix.torch }}, tf_latest)"
-  multi-transformers:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
+        run: pip install nox uv && nox -s "${{ matrix.nox_session }}"
+  partial-install:
+    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:
-        tf: [min]
+        test-env: [onnx, torch]
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/ubuntu-setup
       - name: Run unit tests
-        run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_min)"
+        run: pip install nox uv && nox -s "partial_unit(subset='${{ matrix.test-env }}')"
   launcher:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
+    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
@@ -124,33 +137,18 @@ jobs:
           uv venv .venv
           uv pip install -e . pytest
           uv run python3 -m pytest -v
-  partial-install:
-    if: github.event_name == 'pull_request'
-    needs: [linux]
-    runs-on: ubuntu-latest
-    timeout-minutes: 30
-    strategy:
-      fail-fast: false
-      matrix:
-        test-env: [onnx, torch]
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/ubuntu-setup
-      - name: Run unit tests
-        run: pip install nox uv && nox -s "partial_unit(subset='${{ matrix.test-env }}')"
   unit-pr-required-check:
     # Run even if some jobs are skipped
     if: ${{ github.event_name == 'pull_request' && always() }}
-    needs: [linux, windows, multi-py, multi-torch, multi-transformers, partial-install, launcher]
+    needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher]
     runs-on: ubuntu-latest
     steps:
       - name: Required unit tests did not succeed
         if: >-
-          ${{ needs.linux.result != 'success' ||
-          needs.windows.result != 'success' ||
-          needs.multi-py.result != 'success' ||
-          needs.multi-torch.result != 'success' ||
-          needs.multi-transformers.result != 'success' ||
-          needs.partial-install.result != 'success' ||
-          needs.launcher.result != 'success' }}
+          ${{ needs.linux.result != 'success' || (needs.check-file-changes.outputs.any_changed == 'true' && (
+            needs.windows.result != 'success' ||
+            needs.multi-version.result != 'success' ||
+            needs.partial-install.result != 'success' ||
+            needs.launcher.result != 'success'
+          )) }}
         run: exit 1

From 55c55265dbb9b470da8f9180b980decaf9a66c77 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 18 Apr 2026 08:10:07 -0700
Subject: [PATCH 9/9] Further unify pr and non-pr workflow jobs

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/_pr_gate.yml         | 38 ++++++++-----
 .github/workflows/example_tests.yml    | 75 +++++++-------------------
 .github/workflows/gpu_tests.yml        | 25 ++++-----
 .github/workflows/regression_tests.yml | 38 +++++--------
 .github/workflows/unit_tests.yml       | 20 ++++---
 5 files changed, 78 insertions(+), 118 deletions(-)

diff --git a/.github/workflows/_pr_gate.yml b/.github/workflows/_pr_gate.yml
index 28afff3f64d..d1e6dad0f9f 100644
--- a/.github/workflows/_pr_gate.yml
+++ b/.github/workflows/_pr_gate.yml
@@ -14,34 +14,48 @@ on:
 
 jobs:
   check-file-changes:
-    if: startsWith(github.ref, 'refs/heads/pull-request/')
     runs-on: ubuntu-latest
     outputs:
-      any_changed: ${{ steps.changed-tests.outputs.any_changed }}
+      any_changed: ${{ steps.changed-tests.outputs.any_changed || steps.non-pr.outputs.any_changed }}
     steps:
-      - uses: actions/checkout@v6
+      # For non-PR triggers (schedule, workflow_dispatch), always run tests
+      - id: non-pr
+        if: "!startsWith(github.ref, 'refs/heads/pull-request/')"
+        run: echo "any_changed=true" >> $GITHUB_OUTPUT
+      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
-      - id: get-pr-info
+      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+        id: get-pr-info
         uses: nv-gha-runners/get-pr-info@main
-      # Get commit from main branch that is present in the PR to use as base for changed files
-      - id: calculate-merge-base
+      # Extract SHAs from pr-info JSON via shell to avoid fromJSON on potentially-empty outputs
+      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+        id: pr-shas
         env:
-          PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
-          BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }}
+          PR_INFO: ${{ steps.get-pr-info.outputs.pr-info }}
+        run: |
+          echo "head_sha=$(echo "$PR_INFO" | jq -r '.head.sha')" >> $GITHUB_OUTPUT
+          echo "base_sha=$(echo "$PR_INFO" | jq -r '.base.sha')" >> $GITHUB_OUTPUT
+      # Get commit from main branch that is present in the PR to use as base for changed files
+      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+        id: calculate-merge-base
         run: |
-          (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}"
-      - name: Check for changes in test-relevant directories
+          (echo -n "merge-base="; git merge-base "${{ steps.pr-shas.outputs.base_sha }}" "${{ steps.pr-shas.outputs.head_sha }}") | tee --append "${GITHUB_OUTPUT}"
+      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+        name: Check for changes in test-relevant directories
         id: changed-tests
         uses: step-security/changed-files@v46.0.5
         with:
           base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }}
-          sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }}
+          sha: ${{ steps.pr-shas.outputs.head_sha }}
           files: ${{ inputs.files }}
           fail_on_initial_diff_error: true
   wait-checks:
     needs: [check-file-changes]
-    if: needs.check-file-changes.outputs.any_changed == 'true'
+    if: >-
+      startsWith(github.ref, 'refs/heads/pull-request/') &&
+      needs.check-file-changes.outputs.any_changed == 'true'
     uses: ./.github/workflows/_wait_for_checks.yml
     permissions:
       checks: read
diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 55696ea1570..e8d307ccef9 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -30,10 +30,10 @@ jobs:
         tests/examples/**
 
   ##### PyTorch Example Tests (speculative_decoding requires 26.01 image) #####
-  torch-pr:
+  torch:
     needs: [pr-gate]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
-    strategy: &torch_strategy
+    if: needs.pr-gate.outputs.any_changed == 'true'
+    strategy:
       fail-fast: false
       matrix:
         example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity]
@@ -47,21 +47,9 @@ jobs:
       example: ${{ matrix.example }}
       timeout_minutes: 30
       pip_install_extras: "[hf,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-1
-
-  torch-non-pr:
-    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
-    strategy: *torch_strategy
-    uses: ./.github/workflows/_example_tests_runner.yml
-    secrets: inherit
-    with:
-      docker_image: "nvcr.io/nvidia/pytorch:${{ matrix.docker_image || '26.03' }}-py3"
-      example: ${{ matrix.example }}
-      timeout_minutes: 30
-      pip_install_extras: "[hf,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-2
+      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
 
-  ##### TensorRT-LLM Example Tests #####
+  ##### TensorRT-LLM Example Tests (pr/non-pr split: non-pr runs extra autodeploy+eval examples) #####
   trtllm-pr:
     needs: [pr-gate]
     if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
@@ -92,39 +80,23 @@ jobs:
       runner: linux-amd64-gpu-rtxpro6000-latest-2
 
   ##### Megatron Example Tests #####
-  megatron-pr:
+  megatron:
     needs: [pr-gate]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
-    strategy: &nemo_strategy
-      fail-fast: false
-      matrix:
-        example: [megatron_bridge]
-    uses: ./.github/workflows/_example_tests_runner.yml
-    secrets: inherit
-    with:
-      docker_image: "nvcr.io/nvidia/nemo:26.02"
-      example: ${{ matrix.example }}
-      timeout_minutes: 30
-      pip_install_extras: "[hf,puzzletron,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-1
-
-  megatron-non-pr:
-    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
-    strategy: *nemo_strategy
+    if: needs.pr-gate.outputs.any_changed == 'true'
     uses: ./.github/workflows/_example_tests_runner.yml
     secrets: inherit
     with:
       docker_image: "nvcr.io/nvidia/nemo:26.02"
-      example: ${{ matrix.example }}
+      example: megatron_bridge
       timeout_minutes: 30
       pip_install_extras: "[hf,puzzletron,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-2
+      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
 
   ##### ONNX/TensorRT Example Tests #####
-  onnx-pr:
+  onnx:
     needs: [pr-gate]
-    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true'
-    strategy: &onnx_strategy
+    if: needs.pr-gate.outputs.any_changed == 'true'
+    strategy:
       fail-fast: false
       matrix:
         example: [diffusers, torch_onnx]
@@ -134,33 +106,22 @@ jobs:
       docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3"
       example: ${{ matrix.example }}
       pip_install_extras: "[onnx,hf,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-1
-
-  onnx-non-pr:
-    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
-    strategy: *onnx_strategy
-    uses: ./.github/workflows/_example_tests_runner.yml
-    secrets: inherit
-    with:
-      docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3"
-      example: ${{ matrix.example }}
-      pip_install_extras: "[onnx,hf,dev-test]"
-      runner: linux-amd64-gpu-rtxpro6000-latest-2
+      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
 
   ##### Required Check for PR #####
   example-pr-required-check:
     # Run even if example tests are skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [pr-gate, torch-pr, trtllm-pr, megatron-pr, onnx-pr]
+    needs: [pr-gate, torch, trtllm-pr, megatron, onnx]
     runs-on: ubuntu-latest
     steps:
-      - name: Required GPU tests did not succeed
+      - name: Required example tests did not succeed
         if: |
           needs.pr-gate.result != 'success' ||
           (needs.pr-gate.outputs.any_changed == 'true' && (
-            needs.torch-pr.result != 'success' ||
+            needs.torch.result != 'success' ||
             needs.trtllm-pr.result != 'success' ||
-            needs.megatron-pr.result != 'success' ||
-            needs.onnx-pr.result != 'success'
+            needs.megatron.result != 'success' ||
+            needs.onnx.result != 'success'
           ))
         run: exit 1
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 30a872e91b3..628aead7ee8 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -30,10 +30,11 @@ jobs:
         tests/gpu/**
         tests/gpu_megatron/**
         tests/gpu_trtllm/**
-  gpu-tests-pr:
+
+  gpu-tests:
     needs: [pr-gate]
     if: needs.pr-gate.outputs.any_changed == 'true'
-    strategy: &gpu_strategy
+    strategy:
       fail-fast: false
       matrix:
         include:
@@ -47,15 +48,15 @@ jobs:
           - example: gpu_trtllm
             timeout: 30
             container_image: tensorrt-llm/release:1.3.0rc10
-    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
+    runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: ${{ matrix.timeout }}
-    container: &gpu_container
+    container:
       image: nvcr.io/nvidia/${{ matrix.container_image }}
       env:
         GIT_DEPTH: 1000 # For correct version
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
-    steps: &gpu_steps
+    steps:
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
       - name: Setup environment variables
@@ -75,19 +76,13 @@ jobs:
           flags: gpu
           fail_ci_if_error: false # test may be skipped if relevant file changes are not detected
           verbose: true
-  gpu-tests-non-pr:
-    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
-    strategy: *gpu_strategy
-    runs-on: linux-amd64-gpu-rtxpro6000-latest-2
-    timeout-minutes: ${{ matrix.timeout }}
-    container: *gpu_container
-    steps: *gpu_steps
+
   gpu-pr-required-check:
-    # Run even if gpu-tests-pr is skipped
+    # Run even if gpu-tests is skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [pr-gate, gpu-tests-pr]
+    needs: [pr-gate, gpu-tests]
     runs-on: ubuntu-latest
     steps:
       - name: Required GPU tests did not succeed
-        if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests-pr.result != 'success') }}
+        if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests.result != 'success') }}
         run: exit 1
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
index 2635d751b7b..3e0fd6aba6f 100644
--- a/.github/workflows/regression_tests.yml
+++ b/.github/workflows/regression_tests.yml
@@ -32,25 +32,19 @@ jobs:
         examples/dataset/**
         modelopt_recipes/general/speculative_decoding/**
         tools/launcher/**
-  regression-tests-pr:
+
+  regression-tests:
     needs: [pr-gate]
     if: needs.pr-gate.outputs.any_changed == 'true'
-    strategy: &regression_strategy
-      fail-fast: false
-      matrix:
-        include:
-          - example: regression
-            timeout: 15
-            container_image: pytorch:26.01-py3
-    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
-    timeout-minutes: ${{ matrix.timeout }}
-    container: &regression_container
-      image: nvcr.io/nvidia/${{ matrix.container_image }}
+    runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+    timeout-minutes: 15
+    container:
+      image: nvcr.io/nvidia/pytorch:26.01-py3
       env:
         GIT_DEPTH: 1000 # For correct version
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
-    steps: &regression_steps
+    steps:
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
       - name: Setup environment variables
@@ -60,9 +54,7 @@ jobs:
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
-        run: |
-          pip install nox
-          nox -s ${{ matrix.example }}
+        run: python -m pip install nox && nox -s regression
       - name: Upload regression coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
@@ -71,21 +63,15 @@ jobs:
           flags: regression
           fail_ci_if_error: false # test may be skipped if relevant file changes are not detected
           verbose: true
-  regression-tests-non-pr:
-    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
-    strategy: *regression_strategy
-    runs-on: linux-amd64-gpu-rtxpro6000-latest-2
-    timeout-minutes: ${{ matrix.timeout }}
-    container: *regression_container
-    steps: *regression_steps
+
   regression-pr-required-check:
-    # Run even if regression-tests-pr is skipped
+    # Run even if regression-tests is skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [pr-gate, regression-tests-pr]
+    needs: [pr-gate, regression-tests]
     runs-on: ubuntu-latest
     steps:
       - name: Required regression tests did not succeed
         if: |
           needs.pr-gate.result != 'success' ||
-          (needs.pr-gate.outputs.any_changed == 'true' && needs.regression-tests-pr.result != 'success')
+          (needs.pr-gate.outputs.any_changed == 'true' && needs.regression-tests.result != 'success')
         run: exit 1
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
index 889cdbee83f..9832f0cc605 100644
--- a/.github/workflows/unit_tests.yml
+++ b/.github/workflows/unit_tests.yml
@@ -33,15 +33,19 @@ jobs:
     with:
       match_pattern: "^DCO$"
   check-file-changes:
-    if: github.event_name == 'pull_request'
     runs-on: ubuntu-latest
     outputs:
-      any_changed: ${{ steps.changed.outputs.any_changed }}
+      any_changed: ${{ steps.changed.outputs.any_changed || steps.non-pr.outputs.any_changed }}
     steps:
-      - uses: actions/checkout@v6
+      - id: non-pr
+        if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'
+        run: echo "any_changed=true" >> $GITHUB_OUTPUT
+      - if: github.event_name == 'pull_request'
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
-      - name: Check for changes in test-relevant paths
+      - if: github.event_name == 'pull_request'
+        name: Check for changes in test-relevant paths
         id: changed
         uses: step-security/changed-files@v46.0.5
         with:
@@ -72,7 +76,7 @@ jobs:
           fail_ci_if_error: true
           verbose: true
   windows:
-    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    if: needs.check-file-changes.outputs.any_changed == 'true'
     needs: [linux, check-file-changes]
     runs-on: windows-latest
     timeout-minutes: 30
@@ -84,7 +88,7 @@ jobs:
       - name: Run unit tests (without coverage)
         run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)"
   multi-version:
-    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    if: needs.check-file-changes.outputs.any_changed == 'true'
     needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 30
@@ -107,7 +111,7 @@ jobs:
       - name: Run unit tests
         run: pip install nox uv && nox -s "${{ matrix.nox_session }}"
   partial-install:
-    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    if: needs.check-file-changes.outputs.any_changed == 'true'
     needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 30
@@ -121,7 +125,7 @@ jobs:
       - name: Run unit tests
         run: pip install nox uv && nox -s "partial_unit(subset='${{ matrix.test-env }}')"
   launcher:
-    if: github.event_name == 'pull_request' && needs.check-file-changes.outputs.any_changed == 'true'
+    if: needs.check-file-changes.outputs.any_changed == 'true'
     needs: [linux, check-file-changes]
     runs-on: ubuntu-latest
     timeout-minutes: 15