Skip to content

Commit 1e272b4

Browse files
peri044peri044
authored
feat: Cherry pick of Add validators for dynamic shapes in converter registration (#2849)
1 parent a792bb8 commit 1e272b4

14 files changed

+316
-114
lines changed

‎.github/scripts/install-torch-tensorrt-windows.sh‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
1+
#!/usr/bin/env bash
12
set -eou pipefail
2-
source "${BUILD_ENV_FILE}"
3+
# Source conda so it's available to the script environment
4+
source ${BUILD_ENV_FILE}
5+
export EXTRA_INDEX_URL="https://download.pytorch.org/whl/nightly/${CU_VERSION}"
36

4-
# Install test index version of Torch and Torchvision
5-
${CONDA_RUN} ${PIP_INSTALL_TORCH} torchvision
6-
${CONDA_RUN} pip install pyyaml mpmath==1.3.0
7+
# Install all the dependencies required for Torch-TensorRT
8+
${CONDA_RUN} pip install --pre -r ${PWD}/tests/py/requirements.txt --use-deprecated=legacy-resolver --extra-index-url=${EXTRA_INDEX_URL}
79

8-
# Install TRT from PyPi
9-
TRT_VERSION=$(${CONDA_RUN} python -c "import yaml; print(yaml.safe_load(open('dev_dep_versions.yml', 'r'))['__tensorrt_version__'])")
10-
${CONDA_RUN} pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
11-
12-
# Install pre-built Torch-TRT
10+
# Install Torch-TensorRT
1311
${CONDA_RUN} pip install ${RUNNER_ARTIFACT_DIR}/torch_tensorrt*.whl
1412

15-
echo -e "Running test script";
13+
echo -e "Running test script";

‎.github/scripts/install-torch-tensorrt.sh‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,12 @@
22
set -eou pipefail
33
# Source conda so it's available to the script environment
44
source ${BUILD_ENV_FILE}
5-
${CONDA_RUN} ${PIP_INSTALL_TORCH} torchvision
6-
${CONDA_RUN} python -m pip install pyyaml mpmath==1.3.0
7-
export TRT_VERSION=$(${CONDA_RUN} python -c "import versions; versions.tensorrt_version()")
5+
export EXTRA_INDEX_URL="https://download.pytorch.org/whl/nightly/${CU_VERSION}"
6+
7+
# Install all the dependencies required for Torch-TensorRT
8+
${CONDA_RUN} pip install --pre -r ${PWD}/tests/py/requirements.txt --use-deprecated=legacy-resolver --extra-index-url=${EXTRA_INDEX_URL}
89

910
# Install Torch-TensorRT
10-
${CONDA_RUN} python -m pip install /opt/torch-tensorrt-builds/torch_tensorrt*+${CU_VERSION}*.whl tensorrt~=${TRT_VERSION} --extra-index-url=https://pypi.ngc.nvidia.com
11+
${CONDA_RUN} pip install /opt/torch-tensorrt-builds/torch_tensorrt*.whl
1112

12-
echo -e "Running test script";
13+
echo -e "Running test script";

‎.github/workflows/build-test-linux.yml‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,6 @@ jobs:
7979
export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
8080
pushd .
8181
cd tests/modules
82-
${CONDA_RUN} python -m pip install --pre -r ../py/requirements.txt
8382
${CONDA_RUN} python hub.py
8483
popd
8584
pushd .
@@ -113,7 +112,6 @@ jobs:
113112
export USE_HOST_DEPS=1
114113
pushd .
115114
cd tests/py/dynamo
116-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
117115
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 10 conversion/
118116
popd
119117
@@ -141,7 +139,6 @@ jobs:
141139
export USE_HOST_DEPS=1
142140
pushd .
143141
cd tests/py/dynamo
144-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
145142
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_fe_test_results.xml --ir dynamo models/test_models_export.py
146143
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/test_dyn_models.py
147144
popd
@@ -170,7 +167,6 @@ jobs:
170167
export USE_HOST_DEPS=1
171168
pushd .
172169
cd tests/py/dynamo
173-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
174170
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_serde_test_results.xml --ir dynamo models/test_export_serde.py
175171
popd
176172
@@ -198,7 +194,6 @@ jobs:
198194
export USE_HOST_DEPS=1
199195
pushd .
200196
cd tests/py/dynamo
201-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
202197
${CONDA_RUN} python -m pytest -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
203198
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_comple_be_e2e_test_results.xml --ir torch_compile models/test_models.py
204199
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_dyn_models_export.xml --ir torch_compile models/test_dyn_models.py
@@ -228,7 +223,6 @@ jobs:
228223
export USE_HOST_DEPS=1
229224
pushd .
230225
cd tests/py/dynamo
231-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
232226
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml runtime/
233227
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
234228
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
@@ -257,6 +251,5 @@ jobs:
257251
export USE_HOST_DEPS=1
258252
pushd .
259253
cd tests/py/core
260-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
261254
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
262255
popd

‎.github/workflows/build-test-windows.yml‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,6 @@ jobs:
7272
export USE_HOST_DEPS=1
7373
pushd .
7474
cd tests/py/dynamo
75-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
7675
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 10 conversion/
7776
popd
7877
@@ -98,7 +97,6 @@ jobs:
9897
export USE_HOST_DEPS=1
9998
pushd .
10099
cd tests/py/dynamo
101-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
102100
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_fe_test_results.xml --ir dynamo models/test_models_export.py
103101
${CONDA_RUN} python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/test_dyn_models.py
104102
popd
@@ -125,7 +123,6 @@ jobs:
125123
export USE_HOST_DEPS=1
126124
pushd .
127125
cd tests/py/dynamo
128-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
129126
${CONDA_RUN} python -m pytest -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
130127
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_comple_be_e2e_test_results.xml --ir torch_compile models/test_models.py
131128
popd
@@ -152,7 +149,6 @@ jobs:
152149
export USE_HOST_DEPS=1
153150
pushd .
154151
cd tests/py/dynamo
155-
${CONDA_RUN} python -m pip install --pre -r ../requirements.txt --use-deprecated=legacy-resolver
156152
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml runtime/
157153
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
158154
${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/

‎py/torch_tensorrt/dynamo/_compiler.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def compile(
4747
*,
4848
device: Optional[Union[Device, torch.device, str]] = _defaults.DEVICE,
4949
disable_tf32: bool = _defaults.DISABLE_TF32,
50+
assume_dynamic_shape_support: bool = _defaults.ASSUME_DYNAMIC_SHAPE_SUPPORT,
5051
sparse_weights: bool = _defaults.SPARSE_WEIGHTS,
5152
enabled_precisions: (
5253
Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype]
@@ -106,6 +107,7 @@ def compile(
106107
device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)
107108
108109
disable_tf32 (bool): Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas
110+
assume_dynamic_shape_support (bool): Setting this to true enables the converters work for both dynamic and static shapes. Default: False
109111
sparse_weights (bool): Enable sparsity for convolution and fully connected layers.
110112
enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
111113
refit (bool): Enable refitting
@@ -189,6 +191,7 @@ def compile(
189191
),
190192
"debug": debug,
191193
"device": device,
194+
"assume_dynamic_shape_support": assume_dynamic_shape_support,
192195
"workspace_size": workspace_size,
193196
"min_block_size": min_block_size,
194197
"torch_executed_ops": (
@@ -239,6 +242,9 @@ def compile_module(
239242
"""
240243
dryrun_tracker = DryRunTracker()
241244

245+
# Assume converters support dynamic shapes and disable validation
246+
CONVERTERS.set_dynamic_shape_support(settings.assume_dynamic_shape_support)
247+
242248
# Set torch-executed ops
243249
CONVERTERS.set_disallowed_targets(settings.torch_executed_ops)
244250

@@ -443,6 +449,7 @@ def convert_module_to_trt_engine(
443449
Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype]
444450
) = _defaults.ENABLED_PRECISIONS,
445451
debug: bool = _defaults.DEBUG,
452+
assume_dynamic_shape_support: bool = _defaults.ASSUME_DYNAMIC_SHAPE_SUPPORT,
446453
workspace_size: int = _defaults.WORKSPACE_SIZE,
447454
min_block_size: int = _defaults.MIN_BLOCK_SIZE,
448455
torch_executed_ops: Optional[Set[str]] = None,
@@ -550,6 +557,7 @@ def convert_module_to_trt_engine(
550557
enabled_precisions = {dtype._from(e) for e in enabled_precisions}
551558

552559
compilation_options = {
560+
"assume_dynamic_shape_support": assume_dynamic_shape_support,
553561
"enabled_precisions": enabled_precisions,
554562
"debug": debug,
555563
"workspace_size": workspace_size,
@@ -589,6 +597,10 @@ def convert_module_to_trt_engine(
589597

590598
settings = CompilationSettings(**compilation_options)
591599
logger.info("Compilation Settings: %s\n", settings)
600+
601+
# Assume converters support dynamic shapes and disable validation
602+
CONVERTERS.set_dynamic_shape_support(settings.assume_dynamic_shape_support)
603+
592604
try:
593605
interpreter_result = interpret_module_to_result(gm, input_list, settings)
594606
except UnsupportedOperatorException:

‎py/torch_tensorrt/dynamo/_defaults.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
DEBUG = False
77
DEVICE = None
88
DISABLE_TF32 = False
9+
ASSUME_DYNAMIC_SHAPE_SUPPORT = False
910
DLA_LOCAL_DRAM_SIZE = 1073741824
1011
DLA_GLOBAL_DRAM_SIZE = 536870912
1112
DLA_SRAM_SIZE = 1048576

‎py/torch_tensorrt/dynamo/_settings.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from torch_tensorrt._Device import Device
66
from torch_tensorrt._enums import EngineCapability, dtype
77
from torch_tensorrt.dynamo._defaults import (
8+
ASSUME_DYNAMIC_SHAPE_SUPPORT,
89
DEBUG,
910
DISABLE_TF32,
1011
DLA_GLOBAL_DRAM_SIZE,
@@ -57,6 +58,7 @@ class CompilationSettings:
5758
device (Device): GPU to compile the model on
5859
require_full_compilation (bool): Whether to require the graph is fully compiled in TensorRT.
5960
Only applicable for `ir="dynamo"`; has no effect for `torch.compile` path
61+
assume_dynamic_shape_support (bool): Setting this to true enables the converters work for both dynamic and static shapes. Default: False
6062
disable_tf32 (bool): Whether to disable TF32 computation for TRT layers
6163
sparse_weights (bool): Whether to allow the builder to use sparse weights
6264
refit (bool): Whether to build a refittable engine
@@ -87,6 +89,7 @@ class CompilationSettings:
8789
device: Device = field(default_factory=default_device)
8890
require_full_compilation: bool = REQUIRE_FULL_COMPILATION
8991
disable_tf32: bool = DISABLE_TF32
92+
assume_dynamic_shape_support: bool = ASSUME_DYNAMIC_SHAPE_SUPPORT
9093
sparse_weights: bool = SPARSE_WEIGHTS
9194
refit: bool = REFIT
9295
engine_capability: EngineCapability = field(

0 commit comments

Comments
 (0)