From 23ed7ce661cf23df29d3aac3f2c920598c88d1fc Mon Sep 17 00:00:00 2001
From: Mihir Patel <mihir.v.patel7@gmail.com>
Date: Thu, 16 May 2024 13:11:51 -0400
Subject: [PATCH 01/21] add arm

---
 .github/workflows/docker-configure-build-push.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/docker-configure-build-push.yaml b/.github/workflows/docker-configure-build-push.yaml
index 2b6bf4893d..b79dee4ae0 100644
--- a/.github/workflows/docker-configure-build-push.yaml
+++ b/.github/workflows/docker-configure-build-push.yaml
@@ -87,6 +87,7 @@ jobs:
       uses: docker/build-push-action@v3
       with:
         context: ${{ inputs.context }}
+        platforms: linux/amd64,linux/arm64
         tags: ${{ env.IMAGE_TAG }}
         target: ${{ inputs.target }}
         push: ${{ inputs.push }}

From 0f4a0f6edc356a7ba9aceb211d86402f9ad5e8ea Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 12:46:59 -0500
Subject: [PATCH 02/21] Update docker-configure-build-push.yaml

---
 .github/workflows/docker-configure-build-push.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docker-configure-build-push.yaml b/.github/workflows/docker-configure-build-push.yaml
index b79dee4ae0..a1c74940f1 100644
--- a/.github/workflows/docker-configure-build-push.yaml
+++ b/.github/workflows/docker-configure-build-push.yaml
@@ -36,7 +36,7 @@ on:
         required: true
 jobs:
   configure-build-push:
-    runs-on: ubuntu-latest
+    runs-on: mosaic-4wide
     steps:
     - name: Maximize Build Space on Worker
       uses: easimon/maximize-build-space@v4

From e67bc8fc137eef07a6c26ac7893f8a985c5dba4e Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 13:55:50 -0500
Subject: [PATCH 03/21] Update build_matrix.yaml

---
 docker/build_matrix.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index bf743a4ca2..e1eb0713f3 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -236,13 +236,13 @@
   TARGET: composer_stage
   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
+  BASE_IMAGE: ubuntu:22.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu
   MOFED_VERSION: 5.5-1.0.3.2
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.11'
+  PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.0

From 46af3fbbeb65ebac314ffe22185eefdf0f0c8155 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 14:14:41 -0500
Subject: [PATCH 04/21] Update Dockerfile

---
 docker/Dockerfile | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index c147fe45a7..4744a4afe7 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -271,9 +271,9 @@ ARG MOFED_VERSION
 
 RUN if [ -n "$MOFED_VERSION" ] ; then \
         mkdir -p /tmp/mofed && \
-        wget -nv -P /tmp/mofed http://content.mellanox.com/ofed/MLNX_OFED-${MOFED_VERSION}/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-x86_64.tgz && \
-        tar -zxvf /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-x86_64.tgz -C /tmp/mofed && \
-        /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-x86_64/mlnxofedinstall --user-space-only --without-fw-update --force && \
+        wget -nv -P /tmp/mofed http://content.mellanox.com/ofed/MLNX_OFED-${MOFED_VERSION}/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64.tgz && \
+        tar -zxvf /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64.tgz -C /tmp/mofed && \
+        /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64/mlnxofedinstall --user-space-only --without-fw-update --force && \
         rm -rf /tmp/mofed ; \
     fi
 

From 7ac17f31f22b7d35106bfc980bd6821eb1213f86 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 14:30:33 -0500
Subject: [PATCH 05/21] Update build_matrix.yaml

---
 docker/build_matrix.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index e1eb0713f3..866e548690 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -240,7 +240,7 @@
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu
-  MOFED_VERSION: 5.5-1.0.3.2
+  MOFED_VERSION: 5.6-2.0.9.0
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''

From ebcdba7688201ea3da6ee985c366df2d0e79c5ac Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 14:38:42 -0500
Subject: [PATCH 06/21] Update build_matrix.yaml

---
 docker/build_matrix.yaml | 472 +++++++++++++++++++--------------------
 1 file changed, 236 insertions(+), 236 deletions(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index 866e548690..a7c7886ebf 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -1,240 +1,240 @@
 # This file is automatically generated by generate_build_matrix.py. DO NOT EDIT!
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  IMAGE_NAME: torch-2-3-0-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: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.3.0
-  TAGS:
-  - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
-  - mosaicml/pytorch:latest
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.18.0
-- AWS_OFI_NCCL_VERSION: v1.9.1-aws
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  IMAGE_NAME: torch-2-3-0-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.11'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.3.0
-  TAGS:
-  - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws
-  - mosaicml/pytorch:latest-aws
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.18.0
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  CUDA_VERSION: ''
-  IMAGE_NAME: torch-2-3-0-cpu
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.11'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.3.0
-  TAGS:
-  - mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04
-  - mosaicml/pytorch:latest_cpu
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.18.0
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  IMAGE_NAME: torch-2-2-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.11'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.2.2
-  TAGS:
-  - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.17.2
-- AWS_OFI_NCCL_VERSION: v1.9.1-aws
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  IMAGE_NAME: torch-2-2-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.11'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.2.2
-  TAGS:
-  - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.17.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:20.04
-  CUDA_VERSION: ''
-  IMAGE_NAME: torch-2-2-2-cpu
-  MOFED_VERSION: ''
-  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.11'
-  PYTORCH_NIGHTLY_URL: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.2.2
-  TAGS:
-  - mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.17.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  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
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: v1.9.1-aws
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  CUDA_VERSION: 12.1.1
-  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
-  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
-  TARGET: pytorch_stage
-  TORCHVISION_VERSION: 0.16.2
-- AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
-  CUDA_VERSION: 12.1.1
-  IMAGE_NAME: composer-0-22-0
-  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: ''
-  PYTORCH_NIGHTLY_VERSION: ''
-  PYTORCH_VERSION: 2.3.0
-  TAGS:
-  - mosaicml/composer:0.22.0
-  - mosaicml/composer:latest
-  TARGET: composer_stage
-  TORCHVISION_VERSION: 0.18.0
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   IMAGE_NAME: torch-2-3-0-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: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.3.0
+#   TAGS:
+#   - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
+#   - mosaicml/pytorch:latest
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.18.0
+# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   IMAGE_NAME: torch-2-3-0-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.11'
+#   PYTORCH_NIGHTLY_URL: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.3.0
+#   TAGS:
+#   - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws
+#   - mosaicml/pytorch:latest-aws
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.18.0
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: ubuntu:20.04
+#   CUDA_VERSION: ''
+#   IMAGE_NAME: torch-2-3-0-cpu
+#   MOFED_VERSION: ''
+#   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
+#   PYTHON_VERSION: '3.11'
+#   PYTORCH_NIGHTLY_URL: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.3.0
+#   TAGS:
+#   - mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04
+#   - mosaicml/pytorch:latest_cpu
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.18.0
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   IMAGE_NAME: torch-2-2-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.11'
+#   PYTORCH_NIGHTLY_URL: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.2.2
+#   TAGS:
+#   - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.17.2
+# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   IMAGE_NAME: torch-2-2-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.11'
+#   PYTORCH_NIGHTLY_URL: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.2.2
+#   TAGS:
+#   - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.17.2
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: ubuntu:20.04
+#   CUDA_VERSION: ''
+#   IMAGE_NAME: torch-2-2-2-cpu
+#   MOFED_VERSION: ''
+#   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
+#   PYTHON_VERSION: '3.11'
+#   PYTORCH_NIGHTLY_URL: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.2.2
+#   TAGS:
+#   - mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.17.2
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   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
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.16.2
+# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   CUDA_VERSION: 12.1.1
+#   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
+#   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
+#   TARGET: pytorch_stage
+#   TORCHVISION_VERSION: 0.16.2
+# - AWS_OFI_NCCL_VERSION: ''
+#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+#   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
+#   CUDA_VERSION: 12.1.1
+#   IMAGE_NAME: composer-0-22-0
+#   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: ''
+#   PYTORCH_NIGHTLY_VERSION: ''
+#   PYTORCH_VERSION: 2.3.0
+#   TAGS:
+#   - mosaicml/composer:0.22.0
+#   - mosaicml/composer:latest
+#   TARGET: composer_stage
+#   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
   BASE_IMAGE: ubuntu:22.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0

From bc69f9f3988b44fd4625fe448255181ff8f02ef9 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 14:51:49 -0500
Subject: [PATCH 07/21] Update docker-configure-build-push.yaml

---
 .github/workflows/docker-configure-build-push.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docker-configure-build-push.yaml b/.github/workflows/docker-configure-build-push.yaml
index a1c74940f1..61993a435a 100644
--- a/.github/workflows/docker-configure-build-push.yaml
+++ b/.github/workflows/docker-configure-build-push.yaml
@@ -87,7 +87,7 @@ jobs:
       uses: docker/build-push-action@v3
       with:
         context: ${{ inputs.context }}
-        platforms: linux/amd64,linux/arm64
+        platforms: linux/arm64
         tags: ${{ env.IMAGE_TAG }}
         target: ${{ inputs.target }}
         push: ${{ inputs.push }}

From 04fecc05450e8d7e70d713ea1eeea40d3b32d447 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Thu, 16 May 2024 15:13:55 -0500
Subject: [PATCH 08/21] Update build_matrix.yaml

---
 docker/build_matrix.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index a7c7886ebf..967b2ca3cd 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -236,7 +236,7 @@
 #   TARGET: composer_stage
 #   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: ubuntu:22.04
+  BASE_IMAGE: arm64v8/ubuntu:22.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu

From ccfe107a69b8c6e04376e0ebb537237bdaee25f1 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 15:38:11 -0500
Subject: [PATCH 09/21] PIL 9.0.0 compatible with 20.04

---
 docker/Dockerfile        | 4 ++--
 docker/build_matrix.yaml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 4744a4afe7..a9c5cd3b4d 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -272,8 +272,8 @@ ARG MOFED_VERSION
 RUN if [ -n "$MOFED_VERSION" ] ; then \
         mkdir -p /tmp/mofed && \
         wget -nv -P /tmp/mofed http://content.mellanox.com/ofed/MLNX_OFED-${MOFED_VERSION}/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64.tgz && \
-        tar -zxvf /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64.tgz -C /tmp/mofed && \
-        /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64/mlnxofedinstall --user-space-only --without-fw-update --force && \
+        tar -zxvf /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-aarch64.tgz -C /tmp/mofed && \
+        /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-aarch64/mlnxofedinstall --user-space-only --without-fw-update --force && \
         rm -rf /tmp/mofed ; \
     fi
 
diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index 967b2ca3cd..736b96b06a 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -236,11 +236,11 @@
 #   TARGET: composer_stage
 #   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: arm64v8/ubuntu:22.04
+  BASE_IMAGE: arm64v8/ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu
-  MOFED_VERSION: 5.6-2.0.9.0
+  MOFED_VERSION: 5.5-1.0.3.2
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''

From 311ff8a35024ae9c31d8c60c8e99c59888360a6d Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 15:39:07 -0500
Subject: [PATCH 10/21] 3.8

---
 docker/build_matrix.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index 736b96b06a..2b717ab383 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -242,7 +242,7 @@
   IMAGE_NAME: composer-0-22-0-cpu
   MOFED_VERSION: 5.5-1.0.3.2
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.10'
+  PYTHON_VERSION: '3.8'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.0

From a8acb814a065eb1b5f6d41ff05fa635d4344baec Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 16:11:50 -0500
Subject: [PATCH 11/21] SIMD only compat with x86

---
 docker/Dockerfile | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index a9c5cd3b4d..90eef63e2c 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -191,10 +191,13 @@ ARG PILLOW_SIMD_VERSION
 # so when pillow_simd is installed, other packages won't later override it
 COPY pillow_stub /tmp/pillow_stub
 
-RUN pip${PYTHON_VERSION} install --no-cache-dir --upgrade /tmp/pillow_stub && \
-    pip${PYTHON_VERSION} install --no-cache-dir --upgrade pillow_simd==${PILLOW_SIMD_VERSION} && \
-    rm -rf /tmp/pillow_stub
-
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+        pip${PYTHON_VERSION} install --no-cache-dir --upgrade /tmp/pillow_stub && \
+        pip${PYTHON_VERSION} install --no-cache-dir --upgrade pillow_simd==${PILLOW_SIMD_VERSION} && \
+        rm -rf /tmp/pillow_stub; \
+    else \
+        pip${PYTHON_VERSION} install --no-cache-dir --upgrade Pillow==${PILLOW_PSEUDOVERSION}; \
+    fi
 #################
 # Install Pytorch
 #################

From b7ed5437c7196081cae1c36466aca1905ebdc4ec Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 16:39:11 -0500
Subject: [PATCH 12/21] python 3.11

---
 docker/build_matrix.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index 2b717ab383..baa19709bf 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -242,7 +242,7 @@
   IMAGE_NAME: composer-0-22-0-cpu
   MOFED_VERSION: 5.5-1.0.3.2
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-  PYTHON_VERSION: '3.8'
+  PYTHON_VERSION: '3.11'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.0

From 32f39b950f912c6fb3b33faaaf32c08e510d5875 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 17:14:04 -0500
Subject: [PATCH 13/21] try different target

---
 docker/Dockerfile | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 90eef63e2c..f784759218 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -212,10 +212,16 @@ 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} ; \
+        if [ "$TARGETARCH" = "amd64" ]; 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 \
+            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/cpu/fbgemm_gpu-0.5.0%2Bcpu-cp310-cp310-manylinux2014_aarch64.whl \
+                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} \

From 74e6565ce07a5a22ad21e94da4fa11c8e7e2682a Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 17:52:20 -0500
Subject: [PATCH 14/21] remove +cpu from torch arm install

---
 docker/Dockerfile | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index f784759218..a452bbc67d 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -218,10 +218,9 @@ RUN if [ -z "$PYTORCH_NIGHTLY_URL" ] ; then \
                 torch==${PYTORCH_VERSION}+${CUDA_VERSION_TAG} \
                 torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG} ; \
         else \
-            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/cpu/fbgemm_gpu-0.5.0%2Bcpu-cp310-cp310-manylinux2014_aarch64.whl \
-                torch==${PYTORCH_VERSION}+${CUDA_VERSION_TAG} \
-                torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG} ; \
+            pip${PYTHON_VERSION} install --no-cache-dir --find-links https://download.pytorch.org/whl/cpu/torch_stable.html \
+                torch==${PYTORCH_VERSION} \
+                torchvision==${TORCHVISION_VERSION} ; \
     else \
         pip${PYTHON_VERSION} install --no-cache-dir --pre --index-url ${PYTORCH_NIGHTLY_URL} \
             torch==${PYTORCH_VERSION}.${PYTORCH_NIGHTLY_VERSION} \

From 7512400ed25a62a6e9c466097058efa470a0776a Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 18:07:43 -0500
Subject: [PATCH 15/21] missed fi

---
 docker/Dockerfile | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index a452bbc67d..019d0e8dce 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -221,6 +221,7 @@ RUN if [ -z "$PYTORCH_NIGHTLY_URL" ] ; then \
             pip${PYTHON_VERSION} install --no-cache-dir --find-links https://download.pytorch.org/whl/cpu/torch_stable.html \
                 torch==${PYTORCH_VERSION} \
                 torchvision==${TORCHVISION_VERSION} ; \
+        fi ; \
     else \
         pip${PYTHON_VERSION} install --no-cache-dir --pre --index-url ${PYTORCH_NIGHTLY_URL} \
             torch==${PYTORCH_VERSION}.${PYTORCH_NIGHTLY_VERSION} \

From 4eef39741d9333e193a8cc1d2548dbd4001979c1 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Thu, 16 May 2024 18:23:19 -0500
Subject: [PATCH 16/21] ubuntu20.04-aarch64.tgz

---
 docker/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 019d0e8dce..cdd49539a9 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -280,7 +280,7 @@ ARG MOFED_VERSION
 
 RUN if [ -n "$MOFED_VERSION" ] ; then \
         mkdir -p /tmp/mofed && \
-        wget -nv -P /tmp/mofed http://content.mellanox.com/ofed/MLNX_OFED-${MOFED_VERSION}/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu22.04-aarch64.tgz && \
+        wget -nv -P /tmp/mofed http://content.mellanox.com/ofed/MLNX_OFED-${MOFED_VERSION}/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-aarch64.tgz&& \
         tar -zxvf /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-aarch64.tgz -C /tmp/mofed && \
         /tmp/mofed/MLNX_OFED_LINUX-${MOFED_VERSION}-ubuntu20.04-aarch64/mlnxofedinstall --user-space-only --without-fw-update --force && \
         rm -rf /tmp/mofed ; \

From 3e8d61846f6b2b418189946187cd94d1cc777812 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Fri, 17 May 2024 11:39:22 -0500
Subject: [PATCH 17/21] try cuda

---
 docker/build_matrix.yaml | 60 ++++++++++++++++++++--------------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index baa19709bf..aed3ec2e3a 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -206,37 +206,37 @@
 #   - mosaicml/pytorch:2.1.2_cpu-python3.10-ubuntu20.04
 #   TARGET: pytorch_stage
 #   TORCHVISION_VERSION: 0.16.2
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
-#   CUDA_VERSION: 12.1.1
-#   IMAGE_NAME: composer-0-22-0
-#   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: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.3.0
-#   TAGS:
-#   - mosaicml/composer:0.22.0
-#   - mosaicml/composer:latest
-#   TARGET: composer_stage
-#   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: arm64v8/ubuntu:20.04
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
+  CUDA_VERSION: 12.1.1
+  IMAGE_NAME: composer-0-22-0
+  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: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.3.0
+  TAGS:
+  - mosaicml/composer:0.22.0
+  - mosaicml/composer:latest
+  TARGET: composer_stage
+  TORCHVISION_VERSION: 0.18.0
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: arm64/ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu

From b7249a15cecd7c8fb395c60a1136dda14c9ea681 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <kevin.deshawn@databricks.com>
Date: Fri, 17 May 2024 11:48:23 -0500
Subject: [PATCH 18/21] remove arm64v8

---
 docker/build_matrix.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index aed3ec2e3a..c645f578c5 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -236,7 +236,7 @@
   TARGET: composer_stage
   TORCHVISION_VERSION: 0.18.0
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: arm64/ubuntu:20.04
+  BASE_IMAGE: ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0
   CUDA_VERSION: ''
   IMAGE_NAME: composer-0-22-0-cpu

From f73a4af93dd28e90cc23f7035c360222a93b564a Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Fri, 17 May 2024 13:05:06 -0500
Subject: [PATCH 19/21] Update build_matrix.yaml

---
 docker/build_matrix.yaml | 414 +++++++++++++++++++--------------------
 1 file changed, 207 insertions(+), 207 deletions(-)

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index c645f578c5..bf743a4ca2 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -1,211 +1,211 @@
 # This file is automatically generated by generate_build_matrix.py. DO NOT EDIT!
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   IMAGE_NAME: torch-2-3-0-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: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.3.0
-#   TAGS:
-#   - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
-#   - mosaicml/pytorch:latest
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.18.0
-# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   IMAGE_NAME: torch-2-3-0-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.11'
-#   PYTORCH_NIGHTLY_URL: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.3.0
-#   TAGS:
-#   - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws
-#   - mosaicml/pytorch:latest-aws
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.18.0
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: ubuntu:20.04
-#   CUDA_VERSION: ''
-#   IMAGE_NAME: torch-2-3-0-cpu
-#   MOFED_VERSION: ''
-#   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-#   PYTHON_VERSION: '3.11'
-#   PYTORCH_NIGHTLY_URL: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.3.0
-#   TAGS:
-#   - mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04
-#   - mosaicml/pytorch:latest_cpu
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.18.0
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   IMAGE_NAME: torch-2-2-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.11'
-#   PYTORCH_NIGHTLY_URL: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.2.2
-#   TAGS:
-#   - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.17.2
-# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   IMAGE_NAME: torch-2-2-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.11'
-#   PYTORCH_NIGHTLY_URL: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.2.2
-#   TAGS:
-#   - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.17.2
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: ubuntu:20.04
-#   CUDA_VERSION: ''
-#   IMAGE_NAME: torch-2-2-2-cpu
-#   MOFED_VERSION: ''
-#   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
-#   PYTHON_VERSION: '3.11'
-#   PYTORCH_NIGHTLY_URL: ''
-#   PYTORCH_NIGHTLY_VERSION: ''
-#   PYTORCH_VERSION: 2.2.2
-#   TAGS:
-#   - mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.17.2
-# - AWS_OFI_NCCL_VERSION: ''
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   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
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.16.2
-# - AWS_OFI_NCCL_VERSION: v1.9.1-aws
-#   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
-#   CUDA_VERSION: 12.1.1
-#   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
-#   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
-#   TARGET: pytorch_stage
-#   TORCHVISION_VERSION: 0.16.2
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  IMAGE_NAME: torch-2-3-0-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: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.3.0
+  TAGS:
+  - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
+  - mosaicml/pytorch:latest
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.18.0
+- AWS_OFI_NCCL_VERSION: v1.9.1-aws
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  IMAGE_NAME: torch-2-3-0-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.11'
+  PYTORCH_NIGHTLY_URL: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.3.0
+  TAGS:
+  - mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws
+  - mosaicml/pytorch:latest-aws
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.18.0
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: ubuntu:20.04
+  CUDA_VERSION: ''
+  IMAGE_NAME: torch-2-3-0-cpu
+  MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
+  PYTHON_VERSION: '3.11'
+  PYTORCH_NIGHTLY_URL: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.3.0
+  TAGS:
+  - mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04
+  - mosaicml/pytorch:latest_cpu
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.18.0
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  IMAGE_NAME: torch-2-2-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.11'
+  PYTORCH_NIGHTLY_URL: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.2.2
+  TAGS:
+  - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.17.2
+- AWS_OFI_NCCL_VERSION: v1.9.1-aws
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  IMAGE_NAME: torch-2-2-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.11'
+  PYTORCH_NIGHTLY_URL: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.2.2
+  TAGS:
+  - mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.17.2
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: ubuntu:20.04
+  CUDA_VERSION: ''
+  IMAGE_NAME: torch-2-2-2-cpu
+  MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
+  PYTHON_VERSION: '3.11'
+  PYTORCH_NIGHTLY_URL: ''
+  PYTORCH_NIGHTLY_VERSION: ''
+  PYTORCH_VERSION: 2.2.2
+  TAGS:
+  - mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.17.2
+- AWS_OFI_NCCL_VERSION: ''
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  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
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.16.2
+- AWS_OFI_NCCL_VERSION: v1.9.1-aws
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  CUDA_VERSION: 12.1.1
+  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
+  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
+  TARGET: pytorch_stage
+  TORCHVISION_VERSION: 0.16.2
 - AWS_OFI_NCCL_VERSION: ''
   BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.22.0

From 4763e0d1d130851e00bbe8b5413f82e96b022a36 Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Sat, 18 May 2024 03:03:30 -0500
Subject: [PATCH 20/21] test arm64

---
 docker/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index cdd49539a9..a9f92a887e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -93,7 +93,7 @@ RUN if [ -n "$CUDA_VERSION" ] ; then \
         rm -rf /var/lib/apt/lists/* \
         apt-key del 7fa2af80 && \
         mkdir -p /tmp/cuda-keyring && \
-        wget -P /tmp/cuda-keyring https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \
+        wget -P /tmp/cuda-keyring https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/arm64/cuda-keyring_1.0-1_all.deb && \
         dpkg -i /tmp/cuda-keyring/cuda-keyring_1.0-1_all.deb && \
         rm -rf /tmp/cuda-keyring ; \
     fi

From fb9e44042df933e4a7e2770ee26ffe122ceccc0e Mon Sep 17 00:00:00 2001
From: Kevin DeShawn <126115026+KevDevSha@users.noreply.github.com>
Date: Sat, 18 May 2024 03:37:09 -0500
Subject: [PATCH 21/21] Update docker-configure-build-push.yaml

---
 .github/workflows/docker-configure-build-push.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docker-configure-build-push.yaml b/.github/workflows/docker-configure-build-push.yaml
index 61993a435a..819650d51a 100644
--- a/.github/workflows/docker-configure-build-push.yaml
+++ b/.github/workflows/docker-configure-build-push.yaml
@@ -36,7 +36,7 @@ on:
         required: true
 jobs:
   configure-build-push:
-    runs-on: mosaic-4wide
+    runs-on: mosaic-8wide
     steps:
     - name: Maximize Build Space on Worker
       uses: easimon/maximize-build-space@v4