update nightly to torch 2.3 (#2842)

* update nightly to torch 2.3 * tighten --------- Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
mosaicml · Jan 13, 2024 · c9f0c21 · c9f0c21
1 parent 56fa4bd
commit c9f0c21
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 12 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -294,7 +294,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
 
 ###############

diff --git a/docker/README.md b/docker/README.md
@@ -31,7 +31,7 @@ To install composer, once inside the image, run `pip install mosaicml`.
 <!-- BEGIN_PYTORCH_BUILD_MATRIX -->
 | Linux Distro   | Flavor   | PyTorch Version   | CUDA Version        | Python Version   | Docker Tags                                                                              |
 |----------------|----------|-------------------|---------------------|------------------|------------------------------------------------------------------------------------------|
-| Ubuntu 20.04   | Base     | 2.2.0             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.2.0_cu121-nightly20231213-python3.10-ubuntu20.04`                    |
+| Ubuntu 20.04   | Base     | 2.3.0             | 12.1.0 (Infiniband) | 3.10             | `mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04`                    |
 | Ubuntu 20.04   | Base     | 2.1.2             | 12.1.0 (Infiniband) | 3.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`       |

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
@@ -167,7 +167,7 @@
 - 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,10 +185,10 @@
     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: ''

diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -231,14 +231,14 @@ def _main():
         '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': _get_cuda_override('12.1.0'),
         'PYTHON_VERSION': '3.10',
-        'PYTORCH_VERSION': '2.2.0',
+        'PYTORCH_VERSION': '2.3.0',
         'PYTORCH_NIGHTLY_URL': 'https://download.pytorch.org/whl/nightly/cu121',
-        'PYTORCH_NIGHTLY_VERSION': 'dev20231213+cu121',
-        'TAGS': ['mosaicml/pytorch:2.2.0_cu121-nightly20231213-python3.10-ubuntu20.04'],
+        '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'
     }

diff --git a/setup.py b/setup.py
@@ -78,8 +78,8 @@ def package_files(prefix: str, directory: str, extension: str):
     'tqdm>=4.62.3,<5',
     'torchmetrics>=0.10.0,<1.1',
     'torch_optimizer>=0.3.0,<0.4',
-    'torchvision>=0.13.1,<0.19',
-    'torch>=1.13.1,<2.2.1',
+    'torchvision>=0.13.1,<0.20',  # TODO: Tighten before release
+    'torch>=1.13.1,<2.3.0',  # TODO: Tighten before release
     'requests>=2.26.0,<3',
     'numpy>=1.21.5,<1.27.0',
     'psutil>=5.8.0,<6',