DO NOT MERGE. PR to create torch2.3 + FA 2.4.2 image

docker/Dockerfile

@@ -200,16 +200,12 @@ ENV PYTORCH_VERSION=${PYTORCH_VERSION}
 ENV PYTORCH_NIGHTLY_URL=${PYTORCH_NIGHTLY_URL}
 ENV PYTORCH_NIGHTLY_VERSION=${PYTORCH_NIGHTLY_VERSION}
 
-RUN if [ -z "$PYTORCH_NIGHTLY_URL" ] ; then \
-      CUDA_VERSION_TAG=$(python${PYTHON_VERSION} -c "print('cu' + ''.join('${CUDA_VERSION}'.split('.')[:2]) if '${CUDA_VERSION}' else 'cpu')") && \
-        pip${PYTHON_VERSION} install --no-cache-dir --find-links https://download.pytorch.org/whl/torch_stable.html \
-            torch==${PYTORCH_VERSION}+${CUDA_VERSION_TAG} \
-            torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG} ; \
-    else \
-        pip${PYTHON_VERSION} install --no-cache-dir --pre --index-url ${PYTORCH_NIGHTLY_URL} \
-            torch==${PYTORCH_VERSION}.${PYTORCH_NIGHTLY_VERSION} \
-	        torchvision==${TORCHVISION_VERSION}.${PYTORCH_NIGHTLY_VERSION} ; \
-    fi
+RUN pip${PYTHON_VERSION} uninstall -y torch && \
+    pip${PYTHON_VERSION} uninstall -y pytorch-triton && \
+    pip${PYTHON_VERSION} uninstall -y torchvision && \
+    pip${PYTHON_VERSION} install https://download.pytorch.org/whl/nightly/pytorch_triton-2.2.0%2Be28a256d71-cp310-cp310-linux_x86_64.whl && \
+    pip${PYTHON_VERSION} install https://download.pytorch.org/whl/nightly/cu121/torch-2.3.0.dev20240110%2Bcu121-cp310-cp310-linux_x86_64.whl && \
+    pip${PYTHON_VERSION} install https://download.pytorch.org/whl/nightly/cu121/torchvision-0.18.0.dev20240110%2Bcu121-cp310-cp310-linux_x86_64.whl
 
 #####################################
 # Install EFA and AWS-OFI-NCCL plugin
@@ -294,7 +290,10 @@ RUN if [[ -n "$CUDA_VERSION" ]] &&  [[ -z "${PYTORCH_NIGHTLY_URL}" ]]; then \
 RUN if [ -n "$CUDA_VERSION" ] ; then \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir ninja==1.11.1 && \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir --force-reinstall packaging==22.0 && \
-        pip${PYTHON_VERSION} install --no-cache-dir flash-attn==1.0.9; \
+        git clone --branch v2.4.2 https://github.com/Dao-AILab/flash-attention.git && \
+        cd flash-attention && \
+        MAX_JOBS=1 python${PYTHON_VERSION} setup.py install && \
+        cd .. ; \
     fi
 
 ###############

docker/README.md

@@ -13,10 +13,8 @@ all dependencies for both NLP and Vision models. They are built on top of the
 `mosaicml/composer:latest` or `mosaicml/composer:latest_cpu`, which will always be up to date.
 
 <!-- BEGIN_COMPOSER_BUILD_MATRIX -->
-| Composer Version   | CUDA Support   | Docker Tag                                                     |
-|--------------------|----------------|----------------------------------------------------------------|
-| 0.17.2             | Yes            | `mosaicml/composer:latest`, `mosaicml/composer:0.17.2`         |
-| 0.17.2             | No             | `mosaicml/composer:latest_cpu`, `mosaicml/composer:0.17.2_cpu` |
+| Composer Version   | CUDA Support   | Docker Tag   |
+|--------------------|----------------|--------------|
 <!-- END_COMPOSER_BUILD_MATRIX -->
 
 **Note**: For a lightweight installation, we recommended using a [MosaicML PyTorch Image](#pytorch-images) and manually
@@ -29,18 +27,9 @@ The base flavor contains PyTorch pre-installed; the vision flavor also includes
 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.2.0             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.2.0_cu121-nightly20231213-python3.10-ubuntu20.04`                    |
-| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04`         |
-| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.0 (EFA)        | 3.10             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws` |
-| Ubuntu 20.04   | Base     | 2.1.2             | cpu                 | 3.10             | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.1.2_cpu-python3.10-ubuntu20.04`       |
-| Ubuntu 20.04   | Base     | 2.0.1             | 11.8.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04`                                    |
-| Ubuntu 20.04   | Base     | 2.0.1             | 11.8.0 (EFA)        | 3.10             | `mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04-aws`                                |
-| Ubuntu 20.04   | Base     | 2.0.1             | cpu                 | 3.10             | `mosaicml/pytorch:2.0.1_cpu-python3.10-ubuntu20.04`                                      |
-| Ubuntu 20.04   | Base     | 1.13.1            | 11.7.1 (Infiniband) | 3.10             | `mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04`                                   |
-| Ubuntu 20.04   | Base     | 1.13.1            | 11.7.1 (EFA)        | 3.10             | `mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04-aws`                               |
-| Ubuntu 20.04   | Base     | 1.13.1            | cpu                 | 3.10             | `mosaicml/pytorch:1.13.1_cpu-python3.10-ubuntu20.04`                                     |
+| Linux Distro   | Flavor   | PyTorch Version   | CUDA Version        | Python Version   | Docker Tags                                                           |
+|----------------|----------|-------------------|---------------------|------------------|-----------------------------------------------------------------------|
+| Ubuntu 20.04   | Base     | 2.3.0             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04` |
 <!-- END_PYTORCH_BUILD_MATRIX -->
 
 **Note**: The `mosaicml/pytorch:latest`, `mosaicml/pytorch:latest_cpu`, and `mosaicml/pytorch:latest-aws`

docker/build_matrix.yaml

@@ -2,172 +2,7 @@
 - AWS_OFI_NCCL_VERSION: ''
   BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
   CUDA_VERSION: 12.1.0
-  IMAGE_NAME: torch-2-1-2-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.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.1.2
-  TAGS:
-  - mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04
-  - mosaicml/pytorch:latest
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: v1.7.4-aws
-  BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.0
-  IMAGE_NAME: torch-2-1-2-cu121-aws
-  MOFED_VERSION: ''
-  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.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.1.2
-  TAGS:
-  - mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws
-  - mosaicml/pytorch:latest-aws
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  CUDA_VERSION: ''
-  IMAGE_NAME: torch-2-1-2-cpu
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.1.2
-  TAGS:
-  - mosaicml/pytorch:2.1.2_cpu-python3.10-ubuntu20.04
-  - mosaicml/pytorch:latest_cpu
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 11.8.0
-  IMAGE_NAME: torch-2-0-1-cu118
-  MOFED_VERSION: 5.5-1.0.3.2
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: cuda>=11.8 brand=tesla,driver>=470,driver<471 brand=tesla,driver>=515,driver<516
-    brand=unknown,driver>=470,driver<471 brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=470,driver<471
-    brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=470,driver<471 brand=nvidiartx,driver>=515,driver<516
-    brand=geforce,driver>=470,driver<471 brand=geforce,driver>=515,driver<516 brand=quadro,driver>=470,driver<471
-    brand=quadro,driver>=515,driver<516 brand=titan,driver>=470,driver<471 brand=titan,driver>=515,driver<516
-    brand=titanrtx,driver>=470,driver<471 brand=titanrtx,driver>=515,driver<516
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.0.1
-  TAGS:
-  - mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.15.2
-- AWS_OFI_NCCL_VERSION: v1.7.4-aws
-  BASE_IMAGE: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 11.8.0
-  IMAGE_NAME: torch-2-0-1-cu118-aws
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: cuda>=11.8 brand=tesla,driver>=470,driver<471 brand=tesla,driver>=515,driver<516
-    brand=unknown,driver>=470,driver<471 brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=470,driver<471
-    brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=470,driver<471 brand=nvidiartx,driver>=515,driver<516
-    brand=geforce,driver>=470,driver<471 brand=geforce,driver>=515,driver<516 brand=quadro,driver>=470,driver<471
-    brand=quadro,driver>=515,driver<516 brand=titan,driver>=470,driver<471 brand=titan,driver>=515,driver<516
-    brand=titanrtx,driver>=470,driver<471 brand=titanrtx,driver>=515,driver<516
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.0.1
-  TAGS:
-  - mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04-aws
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.15.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  CUDA_VERSION: ''
-  IMAGE_NAME: torch-2-0-1-cpu
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.0.1
-  TAGS:
-  - mosaicml/pytorch:2.0.1_cpu-python3.10-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.15.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 11.7.1
-  IMAGE_NAME: torch-1-13-1-cu117
-  MOFED_VERSION: 5.5-1.0.3.2
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 1.13.1
-  TAGS:
-  - mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.14.1
-- AWS_OFI_NCCL_VERSION: v1.7.4-aws
-  BASE_IMAGE: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 11.7.1
-  IMAGE_NAME: torch-1-13-1-cu117-aws
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 1.13.1
-  TAGS:
-  - mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04-aws
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.14.1
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  CUDA_VERSION: ''
-  IMAGE_NAME: torch-1-13-1-cpu
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 1.13.1
-  TAGS:
-  - mosaicml/pytorch:1.13.1_cpu-python3.10-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.14.1
-- 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-2-0-20231213-cu121
+  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
@@ -185,54 +20,9 @@
     brand=quadrortx,driver>=525,driver<526 brand=titan,driver>=525,driver<526 brand=titanrtx,driver>=525,driver<526
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: https://download.pytorch.org/whl/nightly/cu121
-  PYTORCH_NIGHTLY_VERSION: dev20231213+cu121
-  PYTORCH_VERSION: 2.2.0
+  PYTORCH_NIGHTLY_VERSION: dev20240110+cu121
+  PYTORCH_VERSION: 2.3.0
   TAGS:
-  - mosaicml/pytorch:2.2.0_cu121-nightly20231213-python3.10-ubuntu20.04
+  - 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.17.2
-  CUDA_VERSION: 12.1.0
-  IMAGE_NAME: composer-0-17-2
-  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.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.1.2
-  TAGS:
-  - mosaicml/composer:0.17.2
-  - mosaicml/composer:latest
-  TARGET: composer_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.17.2
-  CUDA_VERSION: ''
-  IMAGE_NAME: composer-0-17-2-cpu
-  MOFED_VERSION: 5.5-1.0.3.2
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.1.2
-  TAGS:
-  - mosaicml/composer:0.17.2_cpu
-  - mosaicml/composer:latest_cpu
-  TARGET: composer_stage
-  TORCHVISION_VERSION: 0.16.2

docker/generate_build_matrix.py

@@ -243,6 +243,23 @@ def _main():
         'TORCHVISION_VERSION': '0.18.0'
     }
     pytorch_entries.append(nightly_entry)
+    pytorch_entries = []
+    nightly_entry_2_3 = {
+        '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': _get_cuda_override('12.1.0'),
+        'PYTHON_VERSION': '3.10',
+        'PYTORCH_VERSION': '2.3.0',
+        'PYTORCH_NIGHTLY_URL': 'https://download.pytorch.org/whl/nightly/cu121',
+        'PYTORCH_NIGHTLY_VERSION': 'dev20240110+cu121',
+        'TAGS': ['mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04'],
+        'TARGET': 'pytorch_stage',
+        'TORCHVISION_VERSION': '0.18.0'
+    }
+    pytorch_entries.append(nightly_entry_2_3)
     composer_entries = []
 
     # The `GIT_COMMIT` is a placeholder and Jenkins will substitute it with the actual git commit for the `composer_staging` images
@@ -277,7 +294,7 @@ def _main():
         }
 
         composer_entries.append(entry)
-
+    composer_entries = []
     with open(os.path.join(os.path.dirname(__name__), 'build_matrix.yaml'), 'w+') as f:
         f.write('# This file is automatically generated by generate_build_matrix.py. DO NOT EDIT!\n')
         yaml.safe_dump(pytorch_entries + composer_entries, f)