Skip to content

Commit 2616888

Browse files
kevalmorabia97kevalmorabia97
committed
Add CI tests for partial installs and multi-versions + torch 2.8 fix
Signed-off-by: Keval Morabia <[email protected]>
1 parent f053d84 commit 2616888

File tree

12 files changed

+135
-31
lines changed

12 files changed

+135
-31
lines changed

‎.github/workflows/gpu_tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ jobs:
2121
runs-on: linux-amd64-gpu-h100-latest-1
2222
timeout-minutes: 60
2323
container:
24-
image: nvcr.io/nvidia/pytorch:25.04-py3
24+
image: nvcr.io/nvidia/pytorch:25.06-py3
2525
env:
2626
GIT_DEPTH: 1000 # For correct version for tests/gpu/torch/quantization/plugins/test_megatron.py
2727
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‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Run unit tests with older supported Python and Torch versions
2+
name: Multi Version Unit tests
3+
4+
on:
5+
pull_request:
6+
branches: [main, release/*]
7+
paths:
8+
- ".github/workflows/multi_version_unit_tests.yml"
9+
- "modelopt/**"
10+
- "tests/**"
11+
- "setup.py"
12+
- "tox.ini"
13+
14+
# Cancel previous runs if new commit is pushed to the same PR
15+
concurrency:
16+
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
17+
cancel-in-progress: true
18+
19+
jobs:
20+
multi-py-unit:
21+
runs-on: ubuntu-latest
22+
timeout-minutes: 30
23+
strategy:
24+
matrix:
25+
py: [10, 11]
26+
steps:
27+
- uses: actions/checkout@v4
28+
- uses: actions/setup-python@v5
29+
with:
30+
python-version: "3.${{ matrix.py }}"
31+
- name: Install dependencies
32+
run: pip install tox
33+
- name: Run unit tests
34+
run: tox -e py3${{ matrix.py }}-torch28-unit
35+
multi-torch-unit:
36+
runs-on: ubuntu-latest
37+
timeout-minutes: 30
38+
strategy:
39+
matrix:
40+
torch: [25, 26, 27]
41+
steps:
42+
- uses: actions/checkout@v4
43+
- uses: actions/setup-python@v5
44+
with:
45+
python-version: "3.12"
46+
- name: Install dependencies
47+
run: pip install tox
48+
- name: Run unit tests
49+
run: tox -e py312-torch${{ matrix.torch }}-unit

‎.github/workflows/partial_unit_tests.yml‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
name: Partial Install Unit tests
2+
3+
on:
4+
pull_request:
5+
branches: [main, release/*]
6+
paths:
7+
- ".github/workflows/partial_unit_tests.yml"
8+
- "modelopt/**"
9+
- "tests/**"
10+
- "setup.py"
11+
- "tox.ini"
12+
13+
# Cancel previous runs if new commit is pushed to the same PR
14+
concurrency:
15+
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
16+
cancel-in-progress: true
17+
18+
jobs:
19+
partial-unit:
20+
runs-on: ubuntu-latest
21+
timeout-minutes: 30
22+
strategy:
23+
matrix:
24+
test-env: [onnx, torch]
25+
steps:
26+
- uses: actions/checkout@v4
27+
- uses: actions/setup-python@v5
28+
with:
29+
python-version: "3.12"
30+
- name: Install dependencies
31+
run: pip install tox
32+
- name: Run unit tests
33+
run: tox -e py312-partial-unit-${{ matrix.test-env }}
34+
partial-unit-windows:
35+
runs-on: windows-latest
36+
timeout-minutes: 30
37+
strategy:
38+
matrix:
39+
test-env: [onnx, torch]
40+
steps:
41+
- uses: actions/checkout@v4
42+
- uses: actions/setup-python@v5
43+
with:
44+
python-version: "3.12"
45+
- name: Install dependencies
46+
run: pip install tox
47+
- name: Run unit tests
48+
run: tox -e py312-partial-unit-${{ matrix.test-env }}

‎.github/workflows/unit_tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,4 @@ jobs:
2727
- name: Install dependencies
2828
run: pip install tox
2929
- name: Run unit tests
30-
run: tox -e py312-torch27-unit
30+
run: tox -e py312-torch28-unit

‎CHANGELOG.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ Model Optimizer Changelog (Linux)
88

99
**Deprecations**
1010

11+
- Deprecate ``torch<2.5`` support.
12+
1113
**New Features**
1214

1315
- (Experimental) Add quantization support for custom TensorRT op in ONNX models.

‎modelopt/torch/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@
2222

2323
from . import distill, nas, opt, prune, quantization, sparsity, speculative, utils
2424

25-
if _Version(_torch_version) < _Version("2.5"):
25+
if _Version(_torch_version) < _Version("2.6"):
2626
_warnings.warn(
27-
"nvidia-modelopt will drop torch<2.5 support in a future release.", DeprecationWarning
27+
"nvidia-modelopt will drop torch<2.6 support in a future release.", DeprecationWarning
2828
)
2929

3030
# Since `hf` dependencies are optional and users have pre-installed transformers, we need to ensure

‎modelopt/torch/_deploy/utils/torch_onnx.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -411,15 +411,18 @@ def get_onnx_bytes_and_metadata(
411411
)
412412
with torch.inference_mode(), autocast, quantizer_context:
413413
if not dynamo_export or Version(torch.__version__) >= Version("2.6"):
414+
additional_kwargs = {}
415+
if not dynamo_export and Version(torch.__version__) >= Version("2.8"):
416+
additional_kwargs["dynamic_axes"] = dynamic_axes
414417
torch.onnx.export(
415418
model,
416419
dummy_input,
417420
onnx_save_path,
418421
input_names=input_names,
419422
output_names=output_names,
420423
opset_version=onnx_opset,
421-
dynamic_axes=dynamic_axes,
422424
dynamo=dynamo_export,
425+
**additional_kwargs,
423426
)
424427
else: # torch < 2.6 with dynamo export
425428
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)

‎modelopt/torch/nas/modules/utils.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ def get_sliced_tensor_by_slices(
4040
tensor_sliced = tensor
4141
for i, _ in enumerate(slices):
4242
if sum(not isinstance(s, slice) for s in slices) < 2:
43-
tensor_sliced = tensor_sliced[slices]
43+
tensor_sliced = tensor_sliced[tuple(slices)]
4444
break
45-
tensor_sliced = tensor_sliced[slices[: i + 1]]
45+
tensor_sliced = tensor_sliced[tuple(slices[: i + 1])]
4646
slices[i] = slice(None) # replace with a vanilla slice ("[:]") for next slicing iteration
4747

4848
# return sliced, contiguous tensor

‎setup.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
"pulp",
4444
"regex",
4545
"safetensors",
46-
"torch>=2.4",
46+
"torch>=2.5",
4747
"torchprofile>=0.0.4",
4848
"torchvision",
4949
]
@@ -58,7 +58,7 @@
5858
"onnxconverter-common",
5959
"onnxruntime~=1.22.0 ; platform_machine == 'aarch64' or platform_system == 'Darwin'",
6060
"onnxruntime-gpu~=1.22.0 ; platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'", # noqa: E501
61-
"onnxruntime-gpu==1.20.0; platform_system == 'Windows'",
61+
"onnxruntime-directml==1.20.0; platform_system == 'Windows'",
6262
"onnxscript", # For test_onnx_dynamo_export unit test
6363
"onnxsim ; python_version < '3.12' and platform_machine != 'aarch64'",
6464
"polygraphy>=0.49.22",

‎tests/unit/torch/nas/plugins/test_hf_nas_save_restore.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,16 @@
1616
import os
1717

1818
import pytest
19-
import torch
2019
from _test_utils.opt_utils import apply_mode_with_sampling
2120

2221
pytest.importorskip("transformers")
23-
import transformers
2422
from _test_utils.torch_model.transformers_models import (
2523
create_tiny_bert_dir,
2624
tf_modelopt_state_and_output_tester,
2725
)
2826
from transformers import AutoModelForQuestionAnswering, BertForQuestionAnswering
2927

3028

31-
@pytest.mark.skipif(
32-
transformers.__version__ >= "4.52" and torch.__version__ <= "2.4",
33-
reason="Skip when transformers > 4.52 and torch <= 2.4",
34-
)
3529
def test_pruned_transformers_save_restore(tmp_path):
3630
tiny_bert_dir = create_tiny_bert_dir(tmp_path)
3731
model_ref = BertForQuestionAnswering.from_pretrained(tiny_bert_dir)

0 commit comments

Comments
 (0)