From c9f0c215f9a3bf492effd96620c0ef9062bb37d9 Mon Sep 17 00:00:00 2001 From: Charles Tang Date: Sat, 13 Jan 2024 08:44:14 -0800 Subject: [PATCH] update nightly to torch 2.3 (#2842) * update nightly to torch 2.3 * tighten --------- Co-authored-by: Mihir Patel --- docker/Dockerfile | 5 ++++- docker/README.md | 2 +- docker/build_matrix.yaml | 8 ++++---- docker/generate_build_matrix.py | 8 ++++---- setup.py | 4 ++-- 5 files changed, 15 insertions(+), 12 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index ea72ebc7b4..7d3acd7b2c 100644 --- 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 index b05733bb1b..523982f8dd 100644 --- a/docker/README.md +++ b/docker/README.md @@ -31,7 +31,7 @@ To install composer, once inside the image, run `pip install mosaicml`. | 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 index cd2efc0e19..845ea360c7 100644 --- 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 index 3ae69f6d77..b8aa136109 100644 --- 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 index 7322bdc49e..1ecab42499 100644 --- 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',