feat: Upgrade Pytorch and TensorRT versions (#1759)

Signed-off-by: Dheeraj Peri <[email protected]>
Co-authored-by: gs-olive <[email protected]>

Author: peri044

.circleci/config.yml

@@ -31,12 +31,7 @@ In the case of building on top of a custom base container, you first must determ
 version of the PyTorch C++ ABI. If your source of PyTorch is pytorch.org, likely this is the pre-cxx11-abi in which case you must modify `//docker/dist-build.sh` to not build the
 C++11 ABI version of Torch-TensorRT.
 
-You can then build the container using:
-
-
-```bash
-docker build --build-arg BASE_IMG=<IMAGE> -f docker/Dockerfile -t torch_tensorrt:latest .
-```
+You can then build the container using the build command in the [docker README](docker/README.md#instructions)
 
 If you would like to build outside a docker container, please follow the section [Compiling Torch-TensorRT](#compiling-torch-tensorrt)
 
@@ -121,17 +116,17 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedd
 These are the following dependencies used to verify the testcases. Torch-TensorRT can work with other versions, but the tests are not guaranteed to pass.
 
 - Bazel 5.2.0
-- Libtorch 2.0.0.dev20230103 (built with CUDA 11.7)
-- CUDA 11.7
-- cuDNN 8.5.0
-- TensorRT 8.5.1.7
+- Libtorch 2.0.1 (built with CUDA 11.8)
+- CUDA 11.8
+- cuDNN 8.8.0
+- TensorRT 8.6.0
 
 ## Prebuilt Binaries and Wheel files
 
 Releases: https://github.com/pytorch/TensorRT/releases
 
 ```
-pip install torch-tensorrt==1.2.0 --find-links https://github.com/pytorch/TensorRT/releases/expanded_assets/v1.2.0
+pip install torch-tensorrt
 ```
 
 ## Compiling Torch-TensorRT
@@ -252,7 +247,7 @@ A tarball with the include files and library can then be found in bazel-bin
 ### Running Torch-TensorRT on a JIT Graph
 
 > Make sure to add LibTorch to your LD_LIBRARY_PATH <br>
-> `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/bazel-Torch-TensorRT/external/libtorch/lib`
+> `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/bazel-TensorRT/external/libtorch/lib`
 
 ``` shell
 bazel run //cpp/bin/torchtrtc -- $(realpath <PATH TO GRAPH>) out.ts <input-size>

README.md

@@ -41,32 +41,27 @@ local_repository(
 new_local_repository(
     name = "cuda",
     build_file = "@//third_party/cuda:BUILD",
-    path = "/usr/local/cuda-11.7/",
+    path = "/usr/local/cuda-11.8/",
 )
 
-new_local_repository(
-    name = "cublas",
-    build_file = "@//third_party/cublas:BUILD",
-    path = "/usr",
-)
 #############################################################################################################
 # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
 #############################################################################################################
 
 http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "8b3b48615169c83c1b643c0efade078ea080b1da598e15fcf01bc59421f3095e",
+    sha256 = "c5174f18c0866421a5738d389aaea0c02f32a1a5be5f0747dc8dd0d96034c9b0",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.0.dev20230219%2Bcu117.zip"],
+    urls = ["https://download.pytorch.org/libtorch/test/cu118/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "aa7fd06079d260ff83c344d043fb84fbd9cf831cf375ed8b5a1b62416817af31",
+    sha256 = "cc19b398cf435e0e34f347ef90fc11c2a42703998330a9c4a9fb0d2291737df7",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-shared-with-deps-2.0.0.dev20230219%2Bcu117.zip"],
+    urls = ["https://download.pytorch.org/libtorch/test/cu118/libtorch-shared-with-deps-latest.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
@@ -76,20 +71,20 @@ http_archive(
 http_archive(
     name = "cudnn",
     build_file = "@//third_party/cudnn/archive:BUILD",
-    sha256 = "5454a6fd94f008728caae9adad993c4e85ef36302e26bce43bea7d458a5e7b6d",
-    strip_prefix = "cudnn-linux-x86_64-8.5.0.96_cuda11-archive",
+    sha256 = "36fff137153ef73e6ee10bfb07f4381240a86fb9fb78ce372414b528cbab2293",
+    strip_prefix = "cudnn-linux-x86_64-8.8.0.121_cuda11-archive",
     urls = [
-        "https://developer.nvidia.com/compute/cudnn/secure/8.5.0/local_installers/11.7/cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz",
+        "https://developer.download.nvidia.com/compute/cudnn/secure/8.8.0/local_installers/11.8/cudnn-linux-x86_64-8.8.0.121_cuda11-archive.tar.xz",
     ],
 )
 
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "39cc7f077057d1363794e8ff51c4cf21a5dbeccf1116b0020ba0dae0f3063076",
-    strip_prefix = "TensorRT-8.5.1.7",
+    sha256 = "c1732a1093c57ab79fa0b687f061be369e449c9c17792b660f3663ecd8fa7b63",
+    strip_prefix = "TensorRT-8.6.0.12",
     urls = [
-        "https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.5.1/tars/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.0/tars/TensorRT-8.6.0.12.Linux.x86_64-gnu.cuda-11.8.tar.gz",
     ],
 )
 

WORKSPACE

@@ -150,10 +150,10 @@ TRTEngine::TRTEngine(
 }
 
 TRTEngine::~TRTEngine() {
-  rt.reset();
   trt_engine_profiler.reset();
   exec_ctx.reset();
   cuda_engine.reset();
+  rt.reset();
 }
 
 void TRTEngine::disable_profiling() {

core/runtime/TRTEngine.cpp

@@ -1,5 +1,5 @@
 # Base image starts with CUDA
-ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04
+ARG BASE_IMG=nvidia/cuda:11.8.0-devel-ubuntu22.04
 FROM ${BASE_IMG} as base
 
 ARG TENSORRT_VERSION

docker/Dockerfile

@@ -4,7 +4,7 @@
 
 * The `Dockerfile` currently uses <a href="https://github.com/bazelbuild/bazelisk">Bazelisk</a> to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in <a href="https://github.com/pytorch/TensorRT#dependencies">dependencies</a>.
   * The desired versions of CUDNN and TensorRT must be specified as build-args, with major, minor, and patch versions as in: `--build-arg TENSORRT_VERSION=a.b.c --build-arg CUDNN_VERSION=x.y.z`
-  * [**Optional**] The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04`, though this is optional
+  * [**Optional**] The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.8.0-devel-ubuntu22.04`, though this is optional
   * [**Optional**] Additionally, the desired Python version can be changed by explicitly setting a version, as in `--build-arg PYTHON_VERSION=3.10`, though this is optional as well.
 
 * This `Dockerfile` installs `pre-cxx11-abi` versions of Pytorch and builds Torch-TRT using `pre-cxx11-abi` libtorch as well.
@@ -17,14 +17,14 @@ Note: By default the container uses the `pre-cxx11-abi` version of Torch + Torch
 
 ### Instructions
 
-- The example below uses CUDNN 8.5.0 and TensorRT 8.5.1
+- The example below uses CUDNN 8.8.0 and TensorRT 8.6.0
 - See <a href="https://github.com/pytorch/TensorRT#dependencies">dependencies</a> for a list of current default dependencies.
 
 > From root of Torch-TensorRT repo
 
 Build:
 ```
-DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=8.5.1 --build-arg CUDNN_VERSION=8.5.0 -f docker/Dockerfile -t torch_tensorrt:latest .
+DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=8.6.0 --build-arg CUDNN_VERSION=8.8.0 -f docker/Dockerfile -t torch_tensorrt:latest .
 ```
 
 Run:

docker/README.md

@@ -3,7 +3,7 @@
 # Example usage: docker run -it -v$(pwd)/..:/workspace/TRTorch build_trtorch_wheel /bin/bash /workspace/TRTorch/py/build_whl.sh
 
 export CXX=g++
-export CUDA_HOME=/usr/local/cuda-11.7
+export CUDA_HOME=/usr/local/cuda-11.8
 export PROJECT_DIR=/workspace/project
 
 cp -r $CUDA_HOME /usr/local/cuda
@@ -108,4 +108,4 @@ libtorchtrt_pre_cxx11_abi() {
     CUDNN_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __cudnn_version__;print(__cudnn_version__)")
     TORCH_VERSION=$(${PY_DIR}/bin/python -c "from torch import __version__;print(__version__.split('+')[0])")
     cp ${PROJECT_DIR}/bazel-bin/libtorchtrt.tar.gz ${PROJECT_DIR}/py/wheelhouse/libtorchtrt-${TORCHTRT_VERSION}-pre-cxx11-abi-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
-}
+}

py/ci/build_whl.sh

@@ -1,8 +1,8 @@
 numpy
 packaging
 pybind11==2.6.2
---extra-index-url https://download.pytorch.org/whl/nightly/cu117
-torch==2.1.0.dev20230314+cu117
-torchvision==0.15.0.dev20230219+cu117
+--extra-index-url https://download.pytorch.org/whl/test/cu118
+torch==2.0.1
+torchvision==0.15.2
 --extra-index-url https://pypi.ngc.nvidia.com
-tensorrt==8.5.1.7
+tensorrt==8.6.0

py/requirements.txt

@@ -31,6 +31,8 @@
 
 FX_ONLY = False
 
+LEGACY = False
+
 RELEASE = False
 
 CI_RELEASE = False
@@ -48,6 +50,10 @@ def get_git_revision_short_hash() -> str:
     FX_ONLY = True
     sys.argv.remove("--fx-only")
 
+if "--legacy" in sys.argv:
+    LEGACY = True
+    sys.argv.remove("--legacy")
+
 if "--release" not in sys.argv:
     __version__ = __version__ + "+" + get_git_revision_short_hash()
 else:
@@ -380,7 +386,7 @@ def run(self):
     long_description=long_description,
     ext_modules=ext_modules,
     install_requires=[
-        "torch>=1.13.1",
+        "torch >=2.0.1,<2.1" if not LEGACY else "torch >=1.13.0,<2.0",
     ],
     setup_requires=[],
     cmdclass={
@@ -409,7 +415,7 @@ def run(self):
         "Topic :: Software Development",
         "Topic :: Software Development :: Libraries",
     ],
-    python_requires=">=3.7",
+    python_requires=">=3.8",
     include_package_data=True,
     package_data={
         "torch_tensorrt": [

py/setup.py

@@ -58,7 +58,7 @@ def _find_lib(name, paths):
 
     elif sys.platform.startswith("linux"):
         LINUX_PATHS = [
-            "/usr/local/cuda/lib64",
+            "/usr/local/cuda-11.8/lib64",
         ]
 
         if "LD_LIBRARY_PATH" in os.environ:

py/torch_tensorrt/__init__.py

@@ -8,14 +8,14 @@ FX2TRT is merged as FX module in Torch-TensorRT
 ```
     $ conda create --name python_env python=3.8
     $ conda activate python_env
-    # Recommend to install PyTorch 1.12 and later
-    $ conda install pytorch torchvision torchtext cudatoolkit=11.3 -c pytorch-nightly
+    # Recommend to install PyTorch 2.0 and later
+    $ conda install pytorch torchvision torchtext cudatoolkit=11.8 -c pytorch-nightly
     # Install TensorRT python package
     $ pip3 install nvidia-pyindex
-    $ pip3 install tensorrt==8.5.1.7
+    $ pip3 install tensorrt==8.6.0.12
     $ git clone https://github.com/pytorch/TensorRT.git
     $ cd TensorRT/py && python setup.py install --fx-only && cd ..
-    $ pyton -c "import torch_tensorrt.fx"
+    $ python -c "import torch_tensorrt.fx"
     # Test an example by
     $ python py/torch_tensorrt/fx/example/lower_example.py
 ```

py/torch_tensorrt/fx/README.md

@@ -165,7 +165,6 @@ def replace_aten_op_with_indices(module: torch.fx.GraphModule) -> torch.fx.Graph
             torch.ops.aten.max_pool3d_with_indices.default,
             torch.ops.aten.native_batch_norm.default,
             torch.ops.aten._native_batch_norm_legit.default,
-            torch.ops.aten._native_batch_norm_legit_no_training.default,
         ):
             modified = True
             if len(n.users) != 1:
@@ -186,16 +185,6 @@ def replace_aten_op_with_indices(module: torch.fx.GraphModule) -> torch.fx.Graph
                 new_args = list(n.args)
                 new_args.append(False)
                 new_args = tuple(new_args)
-            elif (
-                n.target == torch.ops.aten._native_batch_norm_legit_no_training.default
-            ):
-                new_op = torch.ops.aten.batch_norm
-                new_args = list(n.args)
-                new_args.append(False)
-                # _native_batch_norm_legit_no_training doesn't take in a training arg (assumed to be false)
-                # but batchnorm takes in a training arg at position 5.
-                new_args.insert(5, False)
-                new_args = tuple(new_args)
 
             getitem_node = next(iter(n.users))
             with module.graph.inserting_after(getitem_node):

py/torch_tensorrt/fx/passes/lower_basic_pass_aten.py

@@ -54,11 +54,6 @@ def forward(self, x):
             apply_passes=[fuse_permute_linear],
         )
 
-    # TODO: The following test has been disabled due to a bug in TRT 8.5.1.7
-    # with self.linear2. Issue : https://github.com/pytorch/TensorRT/issues/1444
-    @unittest.skip(
-        reason="test_multi_fuse_permute_linear has been disabled due to a bug in TRT 8.5.1.7 https://github.com/pytorch/TensorRT/issues/1444"
-    )
     def test_multi_fuse_permute_linear(self):
         """
         Fusion when permute output is shared by multiple linears

py/torch_tensorrt/fx/test/passes/test_fuse_permute_linear_trt.py

@@ -1,4 +1,4 @@
 __version__ = "1.4.0.rc0"
-__cuda_version__ = "11.7"
-__cudnn_version__ = "8.5"
-__tensorrt_version__ = "8.5"
+__cuda_version__ = "11.8"
+__cudnn_version__ = "8.8"
+__tensorrt_version__ = "8.6"

py/versions.py

@@ -9,8 +9,7 @@ requires = [
     "cffi",
     "typing_extensions",
     "future",
-    "nvidia-pyindex",
-    "nvidia-tensorrt==8.4.3.1"
+    "tensorrt >=8.6,<8.7"
 ]
 
 # Use legacy backend to import local packages in setup.py
@@ -20,7 +19,7 @@ requires = [
 [tool.black]
 # Uncomment if pyproject.toml worked fine to ensure consistency with flake8
 # line-length = 120
-target-versions = ["py37", "py38", "py39", "py310"]
+target-versions = ["py38", "py39", "py310"]
 force-exclude = """
 elu_converter/setup.py
 """