Skip to content

Commit 6d865da

Browse files
gs-olivegs-olive
authored
Merge pull request #2433 from gs-olive/perf_and_version_fixes
cherry-pick/fix: Performance benchmarking fixes and Torch version fix
2 parents f1d7771 + 9804bc1 commit 6d865da

24 files changed

+520
-597
lines changed

‎.circleci/config.yml

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ commands:
109109
sudo docker run --rm --runtime=nvidia --gpus all nvidia/cuda:11.6.2-base-ubuntu20.04 nvidia-smi
110110
111111
install-cudnn:
112-
description: "Install CUDNN 8.8.1"
112+
description: "Install CUDNN 8.9.5"
113113
parameters:
114114
os:
115115
type: string
@@ -119,7 +119,7 @@ commands:
119119
default: "x86_64"
120120
cudnn-version:
121121
type: string
122-
default: "8.8.1.3"
122+
default: "8.9.5.30"
123123
cuda-version:
124124
type: string
125125
default: "cuda12.0"
@@ -198,7 +198,7 @@ commands:
198198
default: "cuda12.0"
199199
cudnn-version:
200200
type: string
201-
default: "8.8.1.3"
201+
default: "8.9.5.30"
202202
trt-version-short:
203203
type: string
204204
default: "8.6.1"
@@ -246,7 +246,7 @@ commands:
246246
default: "8.6.1"
247247
cudnn-version-long:
248248
type: string
249-
default: "8.8.1.3"
249+
default: "8.9.5.30"
250250
steps:
251251
- run:
252252
name: Set up python environment
@@ -269,10 +269,10 @@ commands:
269269
default: "0.16.0.dev20230703+cu121"
270270
torch-build-index:
271271
type: string
272-
default: "https://download.pytorch.org/whl/nightly/cu121"
272+
default: "https://download.pytorch.org/whl/cu121"
273273
torchvision-build-index:
274274
type: string
275-
default: "https://download.pytorch.org/whl/nightly/cu121"
275+
default: "https://download.pytorch.org/whl/cu121"
276276
steps:
277277
- run:
278278
name: Install Torch
@@ -1457,10 +1457,10 @@ parameters:
14571457
default: "0.16.0.dev20230703+cu121"
14581458
torch-build-index:
14591459
type: string
1460-
default: "https://download.pytorch.org/whl/nightly/cu121"
1460+
default: "https://download.pytorch.org/whl/cu121"
14611461
cudnn-version:
14621462
type: string
1463-
default: "8.8.1.3"
1463+
default: "8.9.5.30"
14641464
trt-version-short:
14651465
type: string
14661466
default: "8.6.1"
@@ -1483,7 +1483,7 @@ parameters:
14831483
default: "https://download.pytorch.org/whl/cu117"
14841484
cudnn-version-legacy:
14851485
type: string
1486-
default: "8.8.1.3"
1486+
default: "8.9.5.30"
14871487
trt-version-short-legacy:
14881488
type: string
14891489
default: "8.6.1"

‎.github/workflows/docgen.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,14 @@ jobs:
3030
- name: Install base deps
3131
run: |
3232
python3 -m pip install pip --upgrade
33-
python3 -m pip install pyyaml numpy torch --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu121
33+
python3 -m pip install pyyaml numpy torch --pre --extra-index-url https://download.pytorch.org/whl/cu121
3434
./packaging/pre_build_script.sh
3535
- name: Get HEAD SHA
3636
id: vars
3737
run: echo "sha=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
3838
- name: Build Python Package
3939
run: |
40-
python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu121
40+
python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/cu121
4141
- name: Generate New Docs
4242
run: |
4343
cd docsrc

‎.github/workflows/docker_builder.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ on:
66
branches:
77
- main
88
- nightly
9+
- release/2.1
910

1011
# If pushes to main are made in rapid succession,
1112
# cancel existing docker builds and use newer commits

‎README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ These are the following dependencies used to verify the testcases. Torch-TensorR
118118
- Bazel 6.2.1
119119
- Libtorch 2.1.0
120120
- CUDA 12.1
121-
- cuDNN 8.8.1
121+
- cuDNN 8.9.5
122122
- TensorRT 8.6.1
123123

124124
## Prebuilt Binaries and Wheel files

‎WORKSPACE

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,14 +54,14 @@ http_archive(
5454
name = "libtorch",
5555
build_file = "@//third_party/libtorch:BUILD",
5656
strip_prefix = "libtorch",
57-
urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
57+
urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.1.0%2Bcu121.zip"],
5858
)
5959

6060
http_archive(
6161
name = "libtorch_pre_cxx11_abi",
6262
build_file = "@//third_party/libtorch:BUILD",
6363
strip_prefix = "libtorch",
64-
urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-shared-with-deps-latest.zip"],
64+
urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.1.0%2Bcu121.zip"],
6565
)
6666

6767
# Download these tarballs manually from the NVIDIA website
@@ -71,10 +71,10 @@ http_archive(
7171
http_archive(
7272
name = "cudnn",
7373
build_file = "@//third_party/cudnn/archive:BUILD",
74-
sha256 = "79d77a769c7e7175abc7b5c2ed5c494148c0618a864138722c887f95c623777c",
75-
strip_prefix = "cudnn-linux-x86_64-8.8.1.3_cuda12-archive",
74+
sha256 = "2a2eb89a2ab51071151c6082f1e816c702167a711a9372f9f73a7b5c4b06e01a",
75+
strip_prefix = "cudnn-linux-x86_64-8.9.5.30_cuda12-archive",
7676
urls = [
77-
"https://developer.nvidia.com/downloads/compute/cudnn/secure/8.8.1/local_installers/12.0/cudnn-linux-x86_64-8.8.1.3_cuda12-archive.tar.xz",
77+
"https://developer.nvidia.com/downloads/compute/cudnn/secure/8.9.5/local_installers/12.x/cudnn-linux-x86_64-8.9.5.30_cuda12-archive.tar.xz",
7878
],
7979
)
8080

‎dev_dep_versions.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
__cuda_version__: "12.1"
2-
__cudnn_version__: "8.8"
2+
__cudnn_version__: "8.9"
33
__tensorrt_version__: "8.6"

‎docker/README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,14 @@ Note: By default the container uses the `pre-cxx11-abi` version of Torch + Torch
1717

1818
### Instructions
1919

20-
- The example below uses CUDNN 8.8 and TensorRT 8.6
20+
- The example below uses CUDNN 8.9 and TensorRT 8.6
2121
- See <a href="https://github.com/pytorch/TensorRT#dependencies">dependencies</a> for a list of current default dependencies.
2222

2323
> From root of Torch-TensorRT repo
2424
2525
Build:
2626
```
27-
DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=8.6 --build-arg CUDNN_VERSION=8.8 -f docker/Dockerfile -t torch_tensorrt:latest .
27+
DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=8.6 --build-arg CUDNN_VERSION=8.9 -f docker/Dockerfile -t torch_tensorrt:latest .
2828
```
2929

3030
Run:

‎docker/dist-build.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
TOP_DIR=$(cd $(dirname $0); pwd)/..
44

55
if [[ -z "${USE_CXX11}" ]]; then
6-
BUILD_CMD="python -m pip wheel . --extra-index-url https://download.pytorch.org/whl/nightly/cu121 -w dist"
6+
BUILD_CMD="python -m pip wheel . --extra-index-url https://download.pytorch.org/whl/cu121 -w dist"
77
else
8-
BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" --extra-index-url https://download.pytorch.org/whl/nightly/cu121 -w dist"
8+
BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" --extra-index-url https://download.pytorch.org/whl/cu121 -w dist"
99
fi
1010

1111
# TensorRT restricts our pip version

‎docsrc/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ Tutorials
8181
tutorials/_rendered_examples/dynamo/torch_compile_resnet_example
8282
tutorials/_rendered_examples/dynamo/torch_compile_transformers_example
8383
tutorials/_rendered_examples/dynamo/torch_compile_advanced_usage
84+
tutorials/_rendered_examples/dynamo/torch_compile_stable_diffusion
8485

8586
Python API Documenation
8687
------------------------

‎docsrc/tutorials/images/majestic_castle.png

494 KB
Loading

‎examples/dynamo/README.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,4 @@ a number of ways you can leverage this backend to accelerate inference.
99
* :ref:`torch_compile_resnet`: Compiling a ResNet model using the Torch Compile Frontend for ``torch_tensorrt.compile``
1010
* :ref:`torch_compile_transformer`: Compiling a Transformer model using ``torch.compile``
1111
* :ref:`torch_compile_advanced_usage`: Advanced usage including making a custom backend to use directly with the ``torch.compile`` API
12+
* :ref:`torch_compile_stable_diffusion`: Compiling a Stable Diffusion model using ``torch.compile``

‎examples/dynamo/torch_compile_stable_diffusion.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
"""
2+
.. _torch_compile_stable_diffusion:
3+
4+
Torch Compile Stable Diffusion
5+
======================================================
6+
7+
This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a Stable Diffusion model. A sample output is featured below:
8+
9+
.. image:: /tutorials/images/majestic_castle.png
10+
:width: 512px
11+
:height: 512px
12+
:scale: 50 %
13+
:align: right
14+
"""
15+
16+
# %%
17+
# Imports and Model Definition
18+
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
19+
20+
import torch
21+
from diffusers import DiffusionPipeline
22+
23+
import torch_tensorrt
24+
25+
model_id = "CompVis/stable-diffusion-v1-4"
26+
device = "cuda:0"
27+
28+
# Instantiate Stable Diffusion Pipeline with FP16 weights
29+
pipe = DiffusionPipeline.from_pretrained(
30+
model_id, revision="fp16", torch_dtype=torch.float16
31+
)
32+
pipe = pipe.to(device)
33+
34+
backend = "torch_tensorrt"
35+
36+
# Optimize the UNet portion with Torch-TensorRT
37+
pipe.unet = torch.compile(
38+
pipe.unet,
39+
backend=backend,
40+
options={
41+
"truncate_long_and_double": True,
42+
"precision": torch.float16,
43+
},
44+
dynamic=False,
45+
)
46+
47+
# %%
48+
# Inference
49+
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
50+
51+
prompt = "a majestic castle in the clouds"
52+
image = pipe(prompt).images[0]
53+
54+
image.save("images/majestic_castle.png")
55+
image.show()

‎py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@
153153
aten.transpose.int,
154154
aten.tril.default,
155155
aten.triu.default,
156+
aten.unbind,
156157
aten.unfold,
157158
aten.unfold_backward,
158159
aten.unfold_copy,

‎py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,22 +2,11 @@
22
from typing import Any, Sequence
33

44
import torch
5-
from torch_tensorrt._utils import sanitized_torch_version
5+
from torch._inductor.freezing import ConstantFolder, replace_node_with_constant
66
from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
77
clean_up_graph_after_modifications,
88
)
99

10-
from packaging import version
11-
12-
# Modify import location of utilities based on Torch version
13-
if version.parse(sanitized_torch_version()) < version.parse("2.1.1"):
14-
from torch._inductor.freezing import ConstantFolder, replace_node_with_constant
15-
else:
16-
from torch._inductor.constant_folding import (
17-
ConstantFolder,
18-
replace_node_with_constant,
19-
)
20-
2110
logger = logging.getLogger(__name__)
2211

2312

0 commit comments

Comments
 (0)