Bumps cudnn FE to v1.5.2

Lightning-AI · Jun 26, 2024 · 7f9424f · 7f9424f
1 parent 1aaa463
commit 7f9424f
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 12 deletions.
diff --git a/.azure/docker-build.yml b/.azure/docker-build.yml
@@ -40,10 +40,10 @@ jobs:
       #maxParallel: "3"
       matrix:
         # CUDA 12.1
-        "cuda 12.1 | torch 2.3 | cudnn FE v1.5.1":
-          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.3.0", TRITON_VERSION: "2.3.0", CUDNN_FRONTEND_VERSION: "1.5.1" }
-        "cuda 12.1 | torch 2.4 /nightly | cudnn FE v1.5.1":
-          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND_VERSION: "1.5.1" }
+        "cuda 12.1 | torch 2.3 | cudnn FE v1.5.2":
+          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.3.0", TRITON_VERSION: "2.3.0", CUDNN_FRONTEND_VERSION: "1.5.2" }
+        "cuda 12.1 | torch 2.4 /nightly | cudnn FE v1.5.2":
+          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND_VERSION: "1.5.2" }
         #'cuda 12.1': # this version - '8.9.5.29-1+cuda12.1' for 'libcudnn8' was not found
     # how much time to give 'run always even if cancelled tasks' before stopping them
     cancelTimeoutInMinutes: "2"

diff --git a/.azure/gpu-tests.yml b/.azure/gpu-tests.yml
@@ -17,17 +17,17 @@ jobs:
       matrix:
         # CUDA 12.1
         "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.3 | regular":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_2.3.0-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.3.0-dev"
           CUDA_VERSION_MM: "121"
         "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.3 | distributed":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_2.3.0-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.3.0-dev"
           CUDA_VERSION_MM: "121"
           testing: "distributed"
         "ubuntu22.04 | cuda 12.1 | python 3.10 | torch-nightly | regular":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_main-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_main-dev"
           CUDA_VERSION_MM: "121"
         "ubuntu22.04 | cuda 12.1 | python 3.10 | torch-nightly | distributed":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_main-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_main-dev"
           CUDA_VERSION_MM: "121"
           testing: "distributed"
     # how much time to give 'run always even if cancelled tasks' before stopping them

diff --git a/.azure/notebook-runs.yml b/.azure/notebook-runs.yml
@@ -16,10 +16,10 @@ jobs:
     strategy:
       matrix:
         "ubuntu22.04 | cuda 12.1 | torch 2.3":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_2.3.0-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.3.0-dev"
           CUDA_VERSION_MM: "121"
         "ubuntu22.04 | cuda 12.1 | torch-nightly":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.1-py3.10-pt_main-dev"
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_main-dev"
           CUDA_VERSION_MM: "121"
     # how long to run the job before automatically cancelling
     timeoutInMinutes: "45"

diff --git a/dockers/ubuntu-cuda/Dockerfile b/dockers/ubuntu-cuda/Dockerfile
@@ -20,7 +20,7 @@ ARG IMAGE_TYPE="devel"
 FROM nvidia/cuda:${CUDA_VERSION}-${IMAGE_TYPE}-ubuntu${UBUNTU_VERSION}
 
 ARG CUDNN_VERSION="9.1.0.70"
-ARG CUDNN_FRONTEND_VERSION="1.5.1"
+ARG CUDNN_FRONTEND_VERSION="1.5.2"
 ARG PYTHON_VERSION="3.10"
 ARG TORCH_VERSION="2.2.1"
 ARG TRITON_VERSION="2.2.0"

diff --git a/thunder/executors/cudnnex.py b/thunder/executors/cudnnex.py
@@ -32,7 +32,8 @@ def required_cudnn_version() -> LooseVersion:
     # Using 1.3.0+ because it works better with other libraries (e.g. torch) that also build on top of cudnn
     # Using 1.5.0+ because it handles exception with unsupported graphs better
     # Using 1.5.1 because of a compatibility fix
-    return LooseVersion("1.5.1")
+    # Using 1.5.2 to allow stride 0
+    return LooseVersion("1.5.2")
 
 
 def cudnn_available() -> bool: