Skip to content

build base images for torch 2.7.1 #2764

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
winglian merged 7 commits into from
Jun 11, 2025
Merged

build base images for torch 2.7.1 #2764

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e9c7801
build base images for torch 2.7.1
winglian Jun 5, 2025
3aeaa73
fix: update base docker to use torch 2.7.1
NanoCode012 Jun 5, 2025
b87ff18
fix: update doc for main base to use 2.7.1
NanoCode012 Jun 5, 2025
8a1369d
make sure to install fa2 in base uv too
winglian Jun 5, 2025
c63a169
use no build isolation for uv+flashattn
winglian Jun 7, 2025
90b6c18
install psutil also for fa2
winglian Jun 9, 2025
1053786
longer timeout for flash attn build
winglian Jun 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions .github/workflows/base.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ on:
jobs:
build-base:
if: github.repository_owner == 'axolotl-ai-cloud'
timeout-minutes: 480
# this job needs to be run on self-hosted GPU runners...
runs-on: ubuntu-latest-m
strategy:
Expand Down Expand Up @@ -47,14 +48,14 @@ jobs:
cuda_version: 12.6.3
cudnn_version: ""
python_version: "3.11"
pytorch: 2.7.0
pytorch: 2.7.1
torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
dockerfile: "Dockerfile-base"
Comment on lines 50 to 53

@coderabbitai coderabbitai Bot Jun 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Upgrade PyTorch version to 2.7.1 for CUDA 12.6.3
This matrix entry now targets PyTorch 2.7.1 under CUDA 12.6.3 and Python 3.11, matching the PR’s goal. Please verify that docker/Dockerfile-base correctly consumes the PYTORCH_VERSION build-arg so the image actually uses 2.7.1.

🏁 Script executed:

#!/bin/bash
# Verify Dockerfile-base uses the PYTORCH_VERSION build argument
grep -R "ARG PYTORCH_VERSION" docker/Dockerfile-base
grep -R "PYTORCH_VERSION" docker/Dockerfile-base

Length of output: 365

Update Dockerfile-base default PyTorch version to 2.7.1

The workflow matrix now installs PyTorch 2.7.1, but docker/Dockerfile-base still defaults to 2.1.2. Please update the build-arg and any related conditionals so the image actually uses 2.7.1:

• File: docker/Dockerfile-base

  • Change the default ARG: 
    - ARG PYTORCH_VERSION="2.1.2"
+ ARG PYTORCH_VERSION="2.7.1"
  • Update any version checks (e.g. if [ "$PYTORCH_VERSION" = "2.7.0" ]) to 2.7.1 or remove if no longer needed.

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In docker/Dockerfile-base around the ARG PYTORCH_VERSION declaration and any
conditional logic checking the PyTorch version, update the default
PYTORCH_VERSION build-arg from 2.1.2 to 2.7.1. Also, modify or remove any
version-specific conditionals that check for 2.7.0 to reflect 2.7.1 or to
simplify if those checks are no longer necessary. This ensures the Docker image
build uses PyTorch 2.7.1 as specified in the workflow matrix.

- cuda: "128"
cuda_version: 12.6.3
cudnn_version: ""
python_version: "3.11"
pytorch: 2.7.0
pytorch: 2.7.1
torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
dockerfile: "Dockerfile-base"
- cuda: "128"
Expand Down Expand Up @@ -106,6 +107,7 @@ jobs:
TORCH_CUDA_ARCH_LIST=${{ matrix.torch_cuda_arch_list }}
build-base-uv:
if: github.repository_owner == 'axolotl-ai-cloud'
timeout-minutes: 480
runs-on: ubuntu-latest-m
strategy:
fail-fast: false
Expand All @@ -122,7 +124,7 @@ jobs:
cuda_version: 12.8.1
cudnn_version: ""
python_version: "3.11"
pytorch: 2.7.0
pytorch: 2.7.1
torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
dockerfile: "Dockerfile-uv-base"
steps:
Expand Down
2 changes: 1 addition & 1 deletion docker/Dockerfile-base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,6 @@ RUN git lfs install --skip-repo && \
# The base image ships with `pydantic==1.8.2` which is not working
pip3 install -U --no-cache-dir pydantic==1.10.10

RUN if [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \
RUN if [ "$PYTORCH_VERSION" = "2.7.1" ] ; then \
pip3 install flash-attn==2.7.4.post1; \
fi
2 changes: 1 addition & 1 deletion docker/Dockerfile-base-next
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
WORKDIR /workspace

RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
python3 -m pip install --no-cache-dir -U torch==2.7.0 --extra-index-url https://download.pytorch.org/whl/test/cu$CUDA && \
python3 -m pip install --no-cache-dir -U torch==2.7.1 --extra-index-url https://download.pytorch.org/whl/test/cu$CUDA && \
python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"

Expand Down
6 changes: 5 additions & 1 deletion docker/Dockerfile-uv-base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,12 @@ RUN uv venv --no-project --relocatable axolotl-venv

ENV PATH="/workspace/axolotl-venv/bin:${PATH}"

RUN uv pip install packaging setuptools wheel \
RUN uv pip install packaging setuptools wheel psutil \
&& uv pip install torch==${PYTORCH_VERSION} \
&& uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
&& uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
&& uv pip install awscli pydantic

RUN if [ "$PYTORCH_VERSION" = "2.7.1" ] ; then \
uv pip install --no-build-isolation flash-attn==2.7.4.post1; \
fi
6 changes: 3 additions & 3 deletions docs/docker.qmd
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ format:
This section describes the different Docker images that are released by AxolotlAI at [Docker Hub](https://hub.docker.com/u/axolotlai).

::: {.callout-important}
For Blackwell GPUs, please use the tags with Pytorch 2.7.0 and CUDA 12.8.
For Blackwell GPUs, please use the tags with Pytorch 2.7.1 and CUDA 12.8.
:::

## Base
Expand All @@ -32,8 +32,8 @@ main-base-py{python_version}-cu{cuda_version}-{pytorch_version}

Tags examples:

- `main-base-py3.11-cu128-2.7.0`
- `main-base-py3.11-cu126-2.7.0`
- `main-base-py3.11-cu128-2.7.1`
- `main-base-py3.11-cu126-2.7.1`
- `main-base-py3.11-cu124-2.6.0`
- `main-base-py3.11-cu124-2.5.1`

Expand Down
Loading