diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index c3fc9168ee..15cf29f95a 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -18,10 +18,10 @@ jobs: matrix: include: - name: "2.4.0_cu124" - base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04 + base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu22.04 dep_groups: "[all]" - name: "2.4.0_cu124_aws" - base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws + base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu22.04-aws dep_groups: "[all]" steps: diff --git a/.github/workflows/pr-cpu.yaml b/.github/workflows/pr-cpu.yaml index 056b070143..053ab221db 100644 --- a/.github/workflows/pr-cpu.yaml +++ b/.github/workflows/pr-cpu.yaml @@ -23,7 +23,7 @@ jobs: include: - name: "cpu-2.4.0" pip_deps: "[all-cpu]" - container: mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu20.04 + container: mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu22.04 markers: "not gpu" pytest_command: "coverage run -m pytest" steps: diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 15c83035e0..6ad7593b38 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -93,7 +93,7 @@ jobs: ${{ env.AWS_DOCKER_TAG }} ${{ env.AWS_LATEST_TAG }} build-args: | - BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws + BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu22.04-aws BRANCH_NAME=${{ env.BRANCH_NAME }} DEP_GROUPS=[all] KEEP_FOUNDRY=true @@ -108,7 +108,7 @@ jobs: ${{ env.DOCKER_TAG }} ${{ env.LATEST_TAG }} build-args: | - BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04 + BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu22.04 BRANCH_NAME=${{ env.BRANCH_NAME }} DEP_GROUPS=[all] KEEP_FOUNDRY=true diff --git a/README.md b/README.md index bc4eff48fd..9dd7ccbc00 100644 --- a/README.md +++ b/README.md @@ -128,7 +128,7 @@ You can select a specific commit hash such as `mosaicml/llm-foundry:2.4.0_cu124- | Docker Image | Torch Version | Cuda Version | LLM Foundry dependencies installed? | | ------------------------------------------------------ | ------------- | ----------------- | ----------------------------------- | -| `mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04` | 2.4.0 | 12.4 (Infiniband) | No | +| `mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu22.04` | 2.4.0 | 12.4 (Infiniband) | No | | `mosaicml/llm-foundry:2.4.0_cu124-latest` | 2.4.0 | 12.4 (Infiniband) | Yes | | `mosaicml/llm-foundry:2.4.0_cu124_aws-latest` | 2.4.0 | 12.4 (EFA) | Yes | diff --git a/scripts/train/benchmarking/submit_benchmarks.py b/scripts/train/benchmarking/submit_benchmarks.py index 27f5c26c7d..90d051e207 100644 --- a/scripts/train/benchmarking/submit_benchmarks.py +++ b/scripts/train/benchmarking/submit_benchmarks.py @@ -47,7 +47,7 @@ def parse_args(): parser.add_argument( '--image', type=str, - default='mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04', + default='mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu22.04', ) parser.add_argument( '--git_branch', diff --git a/scripts/train/benchmarking/sweep.sh b/scripts/train/benchmarking/sweep.sh index 97372ee6fd..deca78184d 100755 --- a/scripts/train/benchmarking/sweep.sh +++ b/scripts/train/benchmarking/sweep.sh @@ -2,7 +2,7 @@ PROJECT="tput" GIT_COMMIT="v0.0.4" -IMAGE="mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04" +IMAGE="mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu22.04" CLUSTER_40GB= # TODO for PRECISION in fp8 bf16