Add CI tests for partial installs and multi-versions + torch 2.8 fix

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

.github/workflows/gpu_tests.yml

@@ -21,7 +21,7 @@ jobs:
     runs-on: linux-amd64-gpu-h100-latest-1
     timeout-minutes: 60
     container:
-      image: nvcr.io/nvidia/pytorch:25.04-py3
+      image: nvcr.io/nvidia/pytorch:25.06-py3
       env:
         GIT_DEPTH: 1000 # For correct version for tests/gpu/torch/quantization/plugins/test_megatron.py
         LD_LIBRARY_PATH: "/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}" # Add libcudnn*.so and libnv*.so to path.

.github/workflows/multi_version_unit_tests.yml

@@ -0,0 +1,49 @@
+# Run unit tests with older supported Python and Torch versions
+name: Multi version tests
+
+on:
+  pull_request:
+    branches: [main, release/*]
+    paths:
+      - ".github/workflows/multi_version_unit_tests.yml"
+      - "modelopt/**"
+      - "tests/**"
+      - "setup.py"
+      - "tox.ini"
+
+# Cancel previous runs if new commit is pushed to the same PR
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  multi-py-unit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        py: [10, 11]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.${{ matrix.py }}"
+      - name: Install dependencies
+        run: pip install tox
+      - name: Run unit tests
+        run: tox -e py3${{ matrix.py }}-torch28-unit
+  multi-torch-unit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        torch: [25, 26, 27]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install dependencies
+        run: pip install tox
+      - name: Run unit tests
+        run: tox -e py312-torch${{ matrix.torch }}-unit

.github/workflows/partial_unit_tests.yml

@@ -0,0 +1,33 @@
+name: Partial install unit tests
+
+on:
+  pull_request:
+    branches: [main, release/*]
+    paths:
+      - ".github/workflows/partial_unit_tests.yml"
+      - "modelopt/**"
+      - "tests/**"
+      - "setup.py"
+      - "tox.ini"
+
+# Cancel previous runs if new commit is pushed to the same PR
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  partial-unit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        test-env: [onnx, torch]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install dependencies
+        run: pip install tox
+      - name: Run unit tests
+        run: tox -e py312-partial-unit-${{ matrix.test-env }}

.github/workflows/unit_tests.yml

@@ -17,8 +17,11 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
     timeout-minutes: 30
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -27,4 +30,4 @@ jobs:
       - name: Install dependencies
         run: pip install tox
       - name: Run unit tests
-        run: tox -e py312-torch27-unit
+        run: tox -e py312-torch28-unit

CHANGELOG.rst

@@ -8,6 +8,8 @@ Model Optimizer Changelog (Linux)
 
 **Deprecations**
 
+- Deprecate ``torch<2.5`` support.
+
 **New Features**
 
 - (Experimental) Add quantization support for custom TensorRT op in ONNX models.

modelopt/torch/__init__.py

@@ -22,9 +22,9 @@
 
 from . import distill, nas, opt, prune, quantization, sparsity, speculative, utils
 
-if _Version(_torch_version) < _Version("2.5"):
+if _Version(_torch_version) < _Version("2.6"):
     _warnings.warn(
-        "nvidia-modelopt will drop torch<2.5 support in a future release.", DeprecationWarning
+        "nvidia-modelopt will drop torch<2.6 support in a future release.", DeprecationWarning
     )
 
 # Since `hf` dependencies are optional and users have pre-installed transformers, we need to ensure

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -411,15 +411,18 @@ def get_onnx_bytes_and_metadata(
     )
     with torch.inference_mode(), autocast, quantizer_context:
         if not dynamo_export or Version(torch.__version__) >= Version("2.6"):
+            additional_kwargs = {}
+            if not dynamo_export and Version(torch.__version__) >= Version("2.8"):
+                additional_kwargs["dynamic_axes"] = dynamic_axes
             torch.onnx.export(
                 model,
                 dummy_input,
                 onnx_save_path,
                 input_names=input_names,
                 output_names=output_names,
                 opset_version=onnx_opset,
-                dynamic_axes=dynamic_axes,
                 dynamo=dynamo_export,
+                **additional_kwargs,
             )
         else:  # torch < 2.6 with dynamo export
             export_options = torch.onnx.ExportOptions(dynamic_shapes=True)

modelopt/torch/nas/modules/utils.py

@@ -40,9 +40,9 @@ def get_sliced_tensor_by_slices(
     tensor_sliced = tensor
     for i, _ in enumerate(slices):
         if sum(not isinstance(s, slice) for s in slices) < 2:
-            tensor_sliced = tensor_sliced[slices]
+            tensor_sliced = tensor_sliced[tuple(slices)]
             break
-        tensor_sliced = tensor_sliced[slices[: i + 1]]
+        tensor_sliced = tensor_sliced[tuple(slices[: i + 1])]
         slices[i] = slice(None)  # replace with a vanilla slice ("[:]") for next slicing iteration
 
     # return sliced, contiguous tensor

setup.py

@@ -43,7 +43,7 @@
     "pulp",
     "regex",
     "safetensors",
-    "torch>=2.4",
+    "torch>=2.5",
     "torchprofile>=0.0.4",
     "torchvision",
 ]
@@ -58,7 +58,7 @@
         "onnxconverter-common",
         "onnxruntime~=1.22.0 ; platform_machine == 'aarch64' or platform_system == 'Darwin'",
         "onnxruntime-gpu~=1.22.0 ; platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'",  # noqa: E501
-        "onnxruntime-gpu==1.20.0; platform_system == 'Windows'",
+        "onnxruntime-directml==1.20.0; platform_system == 'Windows'",
         "onnxscript",  # For test_onnx_dynamo_export unit test
         "onnxsim ; python_version < '3.12' and platform_machine != 'aarch64'",
         "polygraphy>=0.49.22",

tests/conftest.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import platform
+
 import pytest
 
 
@@ -36,3 +38,9 @@ def pytest_collection_modifyitems(config, items):
         for item in items:
             if "manual" in item.keywords:
                 item.add_marker(skipper)
+
+
+@pytest.fixture
+def skip_on_windows():
+    if platform.system() == "Windows":
+        pytest.skip("Skipping on Windows")