add pytorch latest (mosaicml#3290)

mvpatel2000 · May 15, 2024 · 3614d1d · 3614d1d
1 parent 423706b
commit 3614d1d
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 14 deletions.
diff --git a/docker/README.md b/docker/README.md
@@ -28,17 +28,17 @@ The [`mosaicml/pytorch`](https://hub.docker.com/r/mosaicml/pytorch) images conta
 To install composer, once inside the image, run `pip install mosaicml`.
 
 <!-- BEGIN_PYTORCH_BUILD_MATRIX -->
-| Linux Distro   | Flavor   | PyTorch Version   | CUDA Version        | Python Version   | Docker Tags                                               |
-|----------------|----------|-------------------|---------------------|------------------|-----------------------------------------------------------|
-| Ubuntu 20.04   | Base     | 2.3.0             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04`     |
-| Ubuntu 20.04   | Base     | 2.3.0             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws` |
-| Ubuntu 20.04   | Base     | 2.3.0             | cpu                 | 3.11             | `mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04`       |
-| Ubuntu 20.04   | Base     | 2.2.2             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04`     |
-| Ubuntu 20.04   | Base     | 2.2.2             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws` |
-| Ubuntu 20.04   | Base     | 2.2.2             | cpu                 | 3.11             | `mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04`       |
-| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.1 (Infiniband) | 3.10             | `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04`     |
-| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.1 (EFA)        | 3.10             | `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws` |
-| Ubuntu 20.04   | Base     | 2.1.2             | cpu                 | 3.10             | `mosaicml/pytorch:2.1.2_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.1 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04`         |
+| Ubuntu 20.04   | Base     | 2.3.0             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04-aws` |
+| Ubuntu 20.04   | Base     | 2.3.0             | cpu                 | 3.11             | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04`       |
+| Ubuntu 20.04   | Base     | 2.2.2             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04`                                    |
+| Ubuntu 20.04   | Base     | 2.2.2             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.2.2_cu121-python3.11-ubuntu20.04-aws`                                |
+| Ubuntu 20.04   | Base     | 2.2.2             | cpu                 | 3.11             | `mosaicml/pytorch:2.2.2_cpu-python3.11-ubuntu20.04`                                      |
+| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.1 (Infiniband) | 3.10             | `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04`                                    |
+| Ubuntu 20.04   | Base     | 2.1.2             | 12.1.1 (EFA)        | 3.10             | `mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws`                                |
+| Ubuntu 20.04   | Base     | 2.1.2             | cpu                 | 3.10             | `mosaicml/pytorch:2.1.2_cpu-python3.10-ubuntu20.04`                                      |
 <!-- END_PYTORCH_BUILD_MATRIX -->
 
 **Note**: The `mosaicml/pytorch:latest`, `mosaicml/pytorch:latest_cpu`, and `mosaicml/pytorch:latest-aws`

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
@@ -24,6 +24,7 @@
   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
@@ -51,6 +52,7 @@
   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: ''
@@ -65,6 +67,7 @@
   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: ''
@@ -223,7 +226,7 @@
     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'
+  PYTHON_VERSION: '3.11'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.0
@@ -239,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.11'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.0

diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -18,7 +18,7 @@
 import tabulate
 import yaml
 
-PRODUCTION_PYTHON_VERSION = '3.10'
+PRODUCTION_PYTHON_VERSION = '3.11'
 PRODUCTION_PYTORCH_VERSION = '2.3.0'