Add CUDA 12.8 torchaudio and torchvision x86 builds (#6244)

pytorch/pytorch#145570 cc @atalman @ptrblck @nWEIdia
pytorch · Feb 3, 2025 · 3660993 · 3660993
1 parent d45e9d5
commit 3660993
Show file tree

Hide file tree

Showing 6 changed files with 308 additions and 5 deletions.
diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -27,7 +27,7 @@
     "release": ["3.9", "3.10", "3.11", "3.12", "3.13"],
 }
 CUDA_ARCHES_DICT = {
-    "nightly": ["11.8", "12.4", "12.6"],
+    "nightly": ["11.8", "12.4", "12.6", "12.8"],
     "test": ["11.8", "12.4", "12.6"],
     "release": ["11.8", "12.4", "12.6"],
 }
@@ -37,11 +37,11 @@
     "release": ["6.1", "6.2.4"],
 }
 
-CUDA_CUDDN_VERSIONS = {
+CUDA_CUDNN_VERSIONS = {
     "11.8": {"cuda": "11.8.0", "cudnn": "9"},
-    "12.1": {"cuda": "12.1.1", "cudnn": "9"},
     "12.4": {"cuda": "12.4.1", "cudnn": "9"},
-    "12.6": {"cuda": "12.6.2", "cudnn": "9"},
+    "12.6": {"cuda": "12.6.3", "cudnn": "9"},
+    "12.8": {"cuda": "12.8.0", "cudnn": "9"},
 }
 
 PACKAGE_TYPES = ["wheel", "conda", "libtorch"]
@@ -456,6 +456,9 @@ def generate_wheels_matrix(
             upload_to_base_bucket = "no"
             if os in (LINUX, WINDOWS):
                 arches += CUDA_ARCHES
+            # todo: remove once windows cuda 12.8 binaries are available
+            if channel == NIGHTLY and os != LINUX:
+                arches.remove("12.8")
 
         if with_rocm == ENABLE and os == LINUX:
             arches += ROCM_ARCHES

diff --git a/tools/scripts/generate_docker_release_matrix.py b/tools/scripts/generate_docker_release_matrix.py
@@ -45,7 +45,7 @@ def generate_docker_matrix(
         docker_image_version = f"{prefix}-nightly:{generate_binary_build_matrix.CURRENT_NIGHTLY_VERSION}.dev{datetime.today().strftime('%Y%m%d')}"
 
     for cuda in generate_binary_build_matrix.CUDA_ARCHES_DICT[channel]:
-        version = generate_binary_build_matrix.CUDA_CUDDN_VERSIONS[cuda]
+        version = generate_binary_build_matrix.CUDA_CUDNN_VERSIONS[cuda]
         for image in DOCKER_IMAGE_TYPES:
             ret.append(
                 {

diff --git a/tools/tests/assets/build_matrix_linux_wheel_cuda.json b/tools/tests/assets/build_matrix_linux_wheel_cuda.json
@@ -60,6 +60,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.9",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_9-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.9",
       "gpu_arch_type": "rocm",
@@ -150,6 +165,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.10",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_10-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.10",
       "gpu_arch_type": "rocm",
@@ -240,6 +270,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.11",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_11-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.11",
       "gpu_arch_type": "rocm",
@@ -330,6 +375,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.12",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_12-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.12",
       "gpu_arch_type": "rocm",
@@ -420,6 +480,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.13",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_13-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.13",
       "gpu_arch_type": "rocm",

diff --git a/tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json b/tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json
@@ -60,6 +60,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.9",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_9-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.10",
       "gpu_arch_type": "cpu",
@@ -120,6 +135,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.10",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_10-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.11",
       "gpu_arch_type": "cpu",
@@ -180,6 +210,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.11",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_11-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.12",
       "gpu_arch_type": "cpu",
@@ -240,6 +285,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.12",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_12-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.13",
       "gpu_arch_type": "cpu",
@@ -299,6 +359,21 @@
       "upload_to_base_bucket": "no",
       "stable_version": "2.6.0",
       "use_split_build": false
+    },
+    {
+      "python_version": "3.13",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_13-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
     }
   ]
 }
diff --git a/tools/tests/assets/build_matrix_linux_wheel_nocpu.json b/tools/tests/assets/build_matrix_linux_wheel_nocpu.json
@@ -45,6 +45,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.9",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_9-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.9",
       "gpu_arch_type": "rocm",
@@ -120,6 +135,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.10",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_10-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.10",
       "gpu_arch_type": "rocm",
@@ -195,6 +225,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.11",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_11-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.11",
       "gpu_arch_type": "rocm",
@@ -270,6 +315,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.12",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_12-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.12",
       "gpu_arch_type": "rocm",
@@ -345,6 +405,21 @@
       "stable_version": "2.6.0",
       "use_split_build": false
     },
+    {
+      "python_version": "3.13",
+      "gpu_arch_type": "cuda",
+      "gpu_arch_version": "12.8",
+      "desired_cuda": "cu128",
+      "container_image": "pytorch/manylinux2_28-builder:cuda12.8",
+      "package_type": "manywheel",
+      "build_name": "manywheel-py3_13-cuda12_8",
+      "validation_runner": "linux.g5.4xlarge.nvidia.gpu",
+      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
+      "channel": "nightly",
+      "upload_to_base_bucket": "no",
+      "stable_version": "2.6.0",
+      "use_split_build": false
+    },
     {
       "python_version": "3.13",
       "gpu_arch_type": "rocm",