commit change

mosaicml · Aug 9, 2024 · 1ce6cff · 1ce6cff
1 parent e0b3c63
commit 1ce6cff
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 13 deletions.
diff --git a/docker/README.md b/docker/README.md
@@ -30,8 +30,8 @@ To install composer, once inside the image, run `pip install mosaicml`.
 <!-- BEGIN_PYTORCH_BUILD_MATRIX -->
 | Linux Distro   | Flavor   | PyTorch Version   | CUDA Version        | Python Version   | Docker Tags                                                                              |
 |----------------|----------|-------------------|---------------------|------------------|------------------------------------------------------------------------------------------|
-| Ubuntu 20.04   | Base     | 2.4.0             | 12.4.0 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04`         |
-| Ubuntu 20.04   | Base     | 2.4.0             | 12.4.0 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws` |
+| Ubuntu 20.04   | Base     | 2.4.0             | 12.5.1 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.4.0_cu125-python3.11-ubuntu20.04`         |
+| Ubuntu 20.04   | Base     | 2.4.0             | 12.5.1 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.4.0_cu125-python3.11-ubuntu20.04-aws` |
 | Ubuntu 20.04   | Base     | 2.4.0             | cpu                 | 3.11             | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu20.04`       |
 | Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04`                                    |
 | Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws`                                |

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
@@ -1,31 +1,31 @@
 # This file is automatically generated by generate_build_matrix.py. DO NOT EDIT!
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.4.0-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.4.0
-  IMAGE_NAME: torch-2-4-0-cu124
+  BASE_IMAGE: nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04
+  CUDA_VERSION: 12.5.1
+  IMAGE_NAME: torch-2-4-0-cu125
   MOFED_VERSION: latest-23.10
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.11'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.4.0
   TAGS:
-  - mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.4.0_cu125-python3.11-ubuntu20.04
   - mosaicml/pytorch:latest
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.19.0
 - AWS_OFI_NCCL_VERSION: v1.9.1-aws
-  BASE_IMAGE: nvidia/cuda:12.4.0-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.4.0
-  IMAGE_NAME: torch-2-4-0-cu124-aws
+  BASE_IMAGE: nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04
+  CUDA_VERSION: 12.5.1
+  IMAGE_NAME: torch-2-4-0-cu125-aws
   MOFED_VERSION: ''
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.11'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.4.0
   TAGS:
-  - mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
+  - mosaicml/pytorch:2.4.0_cu125-python3.11-ubuntu20.04-aws
   - mosaicml/pytorch:latest-aws
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.19.0
@@ -181,9 +181,9 @@
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.17.2
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.4.0-cudnn8-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.23.5
-  CUDA_VERSION: 12.4.0
+  CUDA_VERSION: 12.5.1
   IMAGE_NAME: composer-0-23-5
   MOFED_VERSION: latest-23.10
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''

diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -35,6 +35,8 @@ def _get_torchvision_version(pytorch_version: str):
 def _get_base_image(cuda_version: str):
     if not cuda_version:
         return 'ubuntu:20.04'
+    if cuda_version == '12.5.1':
+        return f'nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04'
     return f'nvidia/cuda:{cuda_version}-cudnn8-devel-ubuntu20.04'
 
 
@@ -43,7 +45,7 @@ def _get_cuda_version(pytorch_version: str, use_cuda: bool):
     if not use_cuda:
         return ''
     if pytorch_version == '2.4.0':
-        return '12.4.0'
+        return '12.5.1'
     if pytorch_version == '2.3.1':
         return '12.1.1'
     if pytorch_version == '2.2.2':