type

mosaicml · Jan 24, 2024 · eba1c94 · eba1c94
1 parent 1b16ddd
commit eba1c94
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 2 deletions.
diff --git a/docker/README.md b/docker/README.md
@@ -31,6 +31,7 @@ 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.3.0             | 12.1.0 (Infiniband) | 3.11             | `mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04`                    |
 | Ubuntu 20.04   | Base     | 2.3.0             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04`                    |
 | Ubuntu 20.04   | Base     | 2.1.2             | 12.1.0 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.1.2_cu121-python3.11-ubuntu20.04`         |
 | Ubuntu 20.04   | Base     | 2.1.2             | 12.1.0 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.1.2_cu121-python3.11-ubuntu20.04-aws` |

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
@@ -199,6 +199,33 @@
   - mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.18.0
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.0
+  IMAGE_NAME: torch-nightly-2-3-0-20240110-cu121
+  MOFED_VERSION: 5.5-1.0.3.2
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: cuda>=12.1 brand=tesla,driver>=450,driver<451 brand=tesla,driver>=470,driver<471
+    brand=unknown,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471
+    brand=geforce,driver>=470,driver<471 brand=geforcertx,driver>=470,driver<471 brand=quadro,driver>=470,driver<471
+    brand=quadrortx,driver>=470,driver<471 brand=titan,driver>=470,driver<471 brand=titanrtx,driver>=470,driver<471
+    brand=tesla,driver>=510,driver<511 brand=unknown,driver>=510,driver<511 brand=nvidia,driver>=510,driver<511
+    brand=nvidiartx,driver>=510,driver<511 brand=geforce,driver>=510,driver<511 brand=geforcertx,driver>=510,driver<511
+    brand=quadro,driver>=510,driver<511 brand=quadrortx,driver>=510,driver<511 brand=titan,driver>=510,driver<511
+    brand=titanrtx,driver>=510,driver<511 brand=tesla,driver>=515,driver<516 brand=unknown,driver>=515,driver<516
+    brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=515,driver<516 brand=geforce,driver>=515,driver<516
+    brand=geforcertx,driver>=515,driver<516 brand=quadro,driver>=515,driver<516 brand=quadrortx,driver>=515,driver<516
+    brand=titan,driver>=515,driver<516 brand=titanrtx,driver>=515,driver<516 brand=tesla,driver>=525,driver<526
+    brand=unknown,driver>=525,driver<526 brand=nvidia,driver>=525,driver<526 brand=nvidiartx,driver>=525,driver<526
+    brand=geforce,driver>=525,driver<526 brand=geforcertx,driver>=525,driver<526 brand=quadro,driver>=525,driver<526
+    brand=quadrortx,driver>=525,driver<526 brand=titan,driver>=525,driver<526 brand=titanrtx,driver>=525,driver<526
+  PYTHON_VERSION: '3.11'
+  PYTORCH_NIGHTLY_URL: https://download.pytorch.org/whl/nightly/cu121
+  PYTORCH_NIGHTLY_VERSION: dev20240110+cu121
+  PYTORCH_VERSION: 2.3.0
+  TAGS:
+  - mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
   BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.18.0

diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -248,7 +248,7 @@ def _main():
         'TARGET': 'pytorch_stage',
         'TORCHVISION_VERSION': '0.18.0'
     }
-    pytorch_entries.append(nightly_entry_10)
+    pytorch_entries.append(nightly_entry_310)
 
     nightly_entry_311 = {
         'AWS_OFI_NCCL_VERSION': '',
@@ -265,7 +265,7 @@ def _main():
         'TARGET': 'pytorch_stage',
         'TORCHVISION_VERSION': '0.18.0'
     }
-    pytorch_entries.append(nightly_entry_11)
+    pytorch_entries.append(nightly_entry_311)
 
     composer_entries = []