diff --git a/docker/README.md b/docker/README.md index 47b62dc313..6e6133c81f 100644 --- a/docker/README.md +++ b/docker/README.md @@ -31,6 +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-nightly20231024-python3.10-ubuntu20.04` | | Ubuntu 20.04 | Base | 2.1.0 | 12.1.0 (Infiniband) | 3.10 | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04` | | Ubuntu 20.04 | Base | 2.1.0 | 12.1.0 (EFA) | 3.10 | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws` | | Ubuntu 20.04 | Base | 2.1.0 | cpu | 3.10 | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.1.0_cpu-python3.10-ubuntu20.04` | diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml index 4ad47bc701..f7f039b03a 100644 --- a/docker/build_matrix.yaml +++ b/docker/build_matrix.yaml @@ -119,6 +119,18 @@ - mosaicml/pytorch:1.13.1_cpu-python3.10-ubuntu20.04 TARGET: pytorch_stage TORCHVISION_VERSION: 0.14.1 +- AWS_OFI_NCCL_VERSION: '' + BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04 + CUDA_VERSION: 12.1.0 + IMAGE_NAME: torch-nightly-2-2-0-20231024-cu121 + MOFED_VERSION: 5.5-1.0.3.2 + PYTHON_VERSION: '3.10' + PYTORCH_NIGHTLY_URL: https://download.pytorch.org/whl/nightly/cu121 + PYTORCH_NIGHTLY_VERSION: dev20231024+cu121 + PYTORCH_VERSION: 2.2.0 + TAGS: + - mosaicml/pytorch:2.2.0_cu121-nightly20231024-python3.10-ubuntu20.04 + TARGET: pytorch_stage - AWS_OFI_NCCL_VERSION: '' BASE_IMAGE: nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04 COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.16.4 diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py index d08eaa49b9..928c9210f8 100644 --- a/docker/generate_build_matrix.py +++ b/docker/generate_build_matrix.py @@ -183,7 +183,20 @@ def _main(): entry['AWS_OFI_NCCL_VERSION'] = 'v1.7.3-aws' pytorch_entries.append(entry) - + nightly_entry = { + '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-20231024-cu121', + 'MOFED_VERSION': '5.5-1.0.3.2', + 'PYTHON_VERSION': '3.10', + 'PYTORCH_VERSION': '2.2.0', + 'PYTORCH_NIGHTLY_URL': 'https://download.pytorch.org/whl/nightly/cu121', + 'PYTORCH_NIGHTLY_VERSION': 'dev20231024+cu121', + 'TAGS': ['mosaicml/pytorch:2.2.0_cu121-nightly20231024-python3.10-ubuntu20.04'], + 'TARGET': 'pytorch_stage', + } + pytorch_entries.append(nightly_entry) composer_entries = [] # The `GIT_COMMIT` is a placeholder and Jenkins will substitute it with the actual git commit for the `composer_staging` images