diff --git a/docker/Dockerfile.ascend b/docker/Dockerfile.ascend index 565f0408..d15548b6 100644 --- a/docker/Dockerfile.ascend +++ b/docker/Dockerfile.ascend @@ -9,14 +9,31 @@ ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ SHELL ["/bin/bash", "-c"] # ---------- System dependencies ---------- -RUN rm -f /etc/apt/apt.conf.d/docker-clean && \ - find /etc/apt/apt.conf.d -maxdepth 1 -type f | xargs -r grep -l "APT::Update::Post-Invoke\|docker-clean" | xargs -r rm -f && \ - apt-get update -y && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - gcc g++ cmake ninja-build libnuma-dev libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \ - wget git curl jq vim build-essential ca-certificates && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* +RUN set -eux; \ + . /etc/os-release; \ + case "${ID,,}" in \ + ubuntu) \ + rm -f /etc/apt/apt.conf.d/docker-clean; \ + find /etc/apt/apt.conf.d -maxdepth 1 -type f | xargs -r grep -l "APT::Update::Post-Invoke\|docker-clean" | xargs -r rm -f; \ + apt-get update -y; \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + gcc g++ cmake ninja-build libnuma-dev libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \ + wget git curl jq vim build-essential ca-certificates; \ + apt-get clean; \ + rm -rf /var/lib/apt/lists/* \ + ;; \ + openeuler) \ + yum install -y \ + gcc gcc-c++ cmake ninja-build numactl-devel mesa-libGL glib2 libSM libXext libXrender \ + wget git curl jq vim make ca-certificates; \ + yum clean all; \ + rm -rf /var/cache/yum \ + ;; \ + *) \ + echo "Unsupported base image OS: ${ID}" >&2; \ + exit 1 \ + ;; \ + esac RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ pip config set global.extra-index-url "https://pypi.org/simple" && \ @@ -30,14 +47,14 @@ RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ # ---------- Install vllm + vllm-ascend ---------- RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \ - git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm && \ - git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm-ascend.git + git clone --depth 1 --branch {vllm_git_ref} https://github.com/vllm-project/vllm && \ + git clone --depth 1 --branch {vllm_ascend_git_ref} https://github.com/vllm-project/vllm-ascend.git RUN ARCH=$(uname -m) && \ source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ source /usr/local/Ascend/nnal/atb/set_env.sh && \ # Install torch & torch_npu & torchvision - pip install torch==2.9.0 torch_npu==2.9.0.post2 torchvision==0.24.0 && \ + pip install torch=={torch_version} torch_npu=={torch_npu_version} torchvision=={torchvision_version} && \ # Install vllm cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ # Install vllm-ascend @@ -46,14 +63,13 @@ RUN ARCH=$(uname -m) && \ # ---------- Clone training-side repositories ---------- RUN git clone --depth 1 --branch {megatron_branch} https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \ git clone --depth 1 --branch {mindspeed_branch} https://gitcode.com/Ascend/MindSpeed.git /MindSpeed && \ - GIT_LFS_SKIP_SMUDGE=1 git clone --depth 1 -b {swift_branch} --single-branch https://github.com/modelscope/ms-swift.git /ms-swift && \ - git clone --depth 1 https://github.com/modelscope/mcore-bridge.git /mcore-bridge + GIT_LFS_SKIP_SMUDGE=1 git clone --depth 1 -b {swift_branch} --single-branch https://github.com/modelscope/ms-swift.git /ms-swift # ---------- Install training-side repositories ---------- RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \ cd /MindSpeed && pip install --no-cache-dir -e . && \ - cd /mcore-bridge && pip install --no-cache-dir -e . && \ + pip install --no-cache-dir mcore-bridge -i https://pypi.org/simple/ -U && \ cd /ms-swift && pip install --no-cache-dir -e . # ---------- Pin torch to the correct version + torch_npu ---------- @@ -65,13 +81,13 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ if [ "$ARCH" = "x86_64" ]; then \ pip install --no-cache-dir --force-reinstall --no-deps \ --index-url https://download.pytorch.org/whl/cpu \ - torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \ + torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \ else \ pip install --no-cache-dir --force-reinstall --no-deps \ - torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \ + torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \ fi && \ pip install --no-cache-dir --force-reinstall --no-deps \ - torch_npu==2.9.0.post2 && \ + torch_npu=={torch_npu_version} && \ rm -rf /root/.cache/pip # ---------- Remove CUDA-only dependencies pulled in by vllm (they cause missing libtorch_cuda.so errors on NPU) ---------- @@ -126,36 +142,19 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ pip install --no-cache-dir omegaconf==2.3.0 && \ pip cache purge -# ---------- Reinstall triton-ascend for the selected CANN version ---------- +# ---------- Install training and evaluation dependencies ---------- +RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \ + TORCH_DEVICE_BACKEND_AUTOLOAD=0 pip install --no-cache-dir "deepspeed<0.19" ray liger_kernel pre-commit -U && \ + pip cache purge + +# ---------- Install triton-ascend ---------- RUN set -eux; \ pip uninstall -y triton || true; \ pip uninstall -y triton-ascend || true; \ - case "${CANN_VERSION}" in \ - 8.5.*) \ - pip install --no-cache-dir --force-reinstall triton-ascend==3.2.0; \ - ;; \ - 9.0.0) \ - PY_ABI="cp$(python -c 'import sys; print(f"{sys.version_info.major}{sys.version_info.minor}")')"; \ - case "${PY_ABI}" in \ - cp310|cp311|cp312|cp313) ;; \ - *) echo "Unsupported Python ABI for triton-ascend 3.2.1: ${PY_ABI}" >&2; exit 1 ;; \ - esac; \ - ARCH="$(uname -m)"; \ - case "${ARCH}" in \ - aarch64|x86_64) ;; \ - *) echo "Unsupported architecture for triton-ascend 3.2.1: ${ARCH}" >&2; exit 1 ;; \ - esac; \ - WHEEL_NAME="triton_ascend-3.2.1-${PY_ABI}-${PY_ABI}-manylinux_2_27_${ARCH}.manylinux_2_28_${ARCH}.whl"; \ - WHEEL_PATH="/tmp/${WHEEL_NAME}"; \ - curl -fL "https://gitcode.com/Ascend/triton-ascend/releases/download/v3.2.1/${WHEEL_NAME}" -o "${WHEEL_PATH}"; \ - pip install --no-cache-dir --force-reinstall "${WHEEL_PATH}"; \ - rm -f "${WHEEL_PATH}"; \ - ;; \ - *) \ - echo "Unsupported CANN_VERSION for triton-ascend install: ${CANN_VERSION}" >&2; \ - exit 1; \ - ;; \ - esac + pip install --no-cache-dir --force-reinstall \ + triton-ascend=={triton_ascend_version} \ + --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple RUN echo 'source /usr/local/Ascend/ascend-toolkit/set_env.sh' >> /root/.bashrc && \ echo '[ -f /usr/local/Ascend/nnal/atb/set_env.sh ] && source /usr/local/Ascend/nnal/atb/set_env.sh' >> /root/.bashrc && \ diff --git a/docker/OVERVIEW.ascend.md b/docker/OVERVIEW.ascend.md index eb33c451..26deea00 100644 --- a/docker/OVERVIEW.ascend.md +++ b/docker/OVERVIEW.ascend.md @@ -10,7 +10,8 @@ ms-swift Ascend images provide a ready-to-use ms-swift environment for Huawei As - Build template: `docker/Dockerfile.ascend` - Build entrypoint: `docker/build_image.py --image_type ascend` - Default base image: `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11` -- Default output tag: `${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-` +- Supported base OSes: Ubuntu and openEuler, selected from the CANN base-image tag +- Default output tag: `${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-` - Ascend runtime environment is sourced from `/usr/local/Ascend/ascend-toolkit/set_env.sh` - If available, NNAL/ATB runtime is sourced from `/usr/local/Ascend/nnal/atb/set_env.sh` @@ -22,45 +23,46 @@ The Ascend Dockerfile installs and configures: | --- | --- | | CANN | inherited from the selected `quay.io/ascend/cann` base image | | Python | inherited from the base image tag, for example `py3.11` | -| PyTorch | `torch==2.9.0` | -| torch-npu | `torch_npu==2.9.0.post2` | -| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` | -| vLLM | source install from `vllm-project/vllm`, default branch `v0.18.0` | -| vLLM Ascend | source install from `vllm-project/vllm-ascend`, default branch `v0.18.0` | +| PyTorch | `torch==2.9.0` by default; configurable with `--torch_version` | +| torch-npu | `torch_npu==2.9.0.post2` by default; configurable with `--torch_npu_version` | +| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` by default; pass both explicitly when overriding `--torch_version` | +| vLLM | source install from `vllm-project/vllm`, default `0.18.0`; configurable with `--vllm_version` | +| vLLM Ascend | source install from `vllm-project/vllm-ascend`, default `0.18.0`; configurable with `--vllm_ascend_version` | | Megatron-LM | source checkout, default branch `v0.15.3` | | MindSpeed | source checkout, default branch `core_r0.15.3` | -| mcore-bridge | source checkout from `modelscope/mcore-bridge` | +| mcore-bridge | latest release from PyPI | | ms-swift | source checkout from `modelscope/ms-swift`, default branch `main` | | ModelScope | source checkout from `modelscope/modelscope`, default branch `master` | -| triton-ascend | `3.2.0` for CANN `8.5.*`; local wheel install of `3.2.1` for CANN `9.0.0` | +| triton-ascend | CANN `8.5.*` defaults to `3.2.0`; CANN `9.0.*` defaults to `3.2.1`; configurable with `--triton_ascend_version` and installed from the Triton Ascend PyPI index | ## Supported Tag Format Images built by `docker/build_image.py --image_type ascend` use this tag format: ```text -${DOCKER_REGISTRY}:----- +${DOCKER_REGISTRY}:--torch_npu---- ``` | Field | Example | Description | | --- | --- | --- | | `swift-branch` | `main` | ms-swift branch used during image build | -| `atlas-hardware` | `A2`, `A3`, `300I`, `A5` | Derived from `--soc_version` | -| `python-tag` | `py311` | Derived from `--python_version` | -| `cann-version-tag` | `CANN8.5.1`, `CANN9.0.0` | Parsed from the CANN base image tag | -| `os-tag` | `ubuntu22.04` | Parsed from the CANN base image tag | -| `arch` | `arm`, `x86` | Derived from host architecture or `--arch` | +| `cann-version-tag` | `cann8.5.1`, `cann9.0.0` | Parsed from the CANN base image tag | +| `torch-npu-version` | `2.9.0.post2` | From `--torch_npu_version`; defaults to `2.9.0.post2` | +| `atlas-hardware` | `a2`, `a3`, `300i`, `a5` | Derived from `--soc_version` | +| `os-tag` | `ubuntu22.04`, `openeuler24.03` | Parsed from the CANN base-image tag; prevents tags for different OSes from colliding | +| `python-tag` | `py3.11` | Parsed from the CANN base image tag | +| `arch` | `aarch64`, `x86_64` | Derived from host architecture or `--arch` | Default example on an ARM64 host: ```text -${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm +${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64 ``` A2 / CANN 9.0.0 example: ```text -${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm +${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 ``` ## Build Locally @@ -85,6 +87,39 @@ python docker/build_image.py \ --soc_version ascend910b1 ``` +Build an openEuler image. The system-dependency layer automatically uses `yum`; Ubuntu images continue to use `apt-get`. + +```bash +python docker/build_image.py \ + --image_type ascend \ + --base_image quay.io/ascend/cann:8.5.1-a3-openeuler24.03-py3.11 \ + --soc_version ascend910_9391 +``` + +Override the PyTorch stack. `--torch_version` must match the base version of +`--torch_npu_version`; when overriding PyTorch, pass its matching torchvision +and torchaudio versions explicitly. + +```bash +python docker/build_image.py \ + --image_type ascend \ + --torch_version 2.9.0 \ + --torch_npu_version 2.9.0.post2 \ + --torchvision_version 0.24.0 \ + --torchaudio_version 2.9.0 +``` + +Override the vLLM stack or triton-ascend. The vLLM version arguments select +the matching Git tag, for example `0.18.0` selects `v0.18.0`. + +```bash +python docker/build_image.py \ + --image_type ascend \ + --vllm_version 0.18.0 \ + --vllm_ascend_version 0.18.0 \ + --triton_ascend_version 3.2.1 +``` + Override Megatron or MindSpeed source branches when needed: ```bash @@ -94,11 +129,11 @@ python docker/build_image.py \ --mindspeed_branch core_r0.15.3 ``` -For slow networks, Linux hosts can use Docker host networking after the root `Dockerfile` is generated: +To run the rendered Dockerfile manually, use: ```bash -docker build --network host \ - -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ +docker build \ + -t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \ -f Dockerfile . ``` @@ -119,7 +154,7 @@ docker run --rm -it \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /mnt/workspace:/mnt/workspace \ - ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \ bash ``` @@ -147,7 +182,8 @@ pip show ms-swift modelscope torch-npu triton-ascend ## Notes - CANN, firmware, and driver versions must be compatible with each other. -- CANN `8.5.*` and CANN `9.0.0` use different `triton-ascend` install paths in this Dockerfile. +- Ubuntu base images install system dependencies through `apt-get`; openEuler base images install the corresponding RPM packages through `yum`. +- `triton-ascend` is installed from `https://triton-ascend.osinfra.cn/pypi/simple`; select a version compatible with the chosen CANN, Python, and architecture. - The image is intended for Ascend NPU ms-swift workflows. CUDA-only packages pulled in by dependencies are removed when they conflict with NPU runtime libraries. - Use a fixed image tag for production jobs instead of relying on a moving branch name. diff --git a/docker/OVERVIEW.ascend.zh.md b/docker/OVERVIEW.ascend.zh.md index d71f2066..2bdde812 100644 --- a/docker/OVERVIEW.ascend.zh.md +++ b/docker/OVERVIEW.ascend.zh.md @@ -10,7 +10,8 @@ ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms - 构建模板：`docker/Dockerfile.ascend` - 构建入口：`docker/build_image.py --image_type ascend` - 默认基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11` -- 默认输出 tag：`${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-` +- 支持的基础 OS：Ubuntu 和 openEuler，由 CANN 基础镜像 tag 选择 +- 默认输出 tag：`${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-` - Ascend runtime 环境来自 `/usr/local/Ascend/ascend-toolkit/set_env.sh` - 如果镜像内存在 NNAL/ATB，则会加载 `/usr/local/Ascend/nnal/atb/set_env.sh` @@ -22,45 +23,46 @@ Ascend Dockerfile 会安装和配置： | --- | --- | | CANN | 继承自选定的 `quay.io/ascend/cann` 基础镜像 | | Python | 继承自基础镜像 tag，例如 `py3.11` | -| PyTorch | `torch==2.9.0` | -| torch-npu | `torch_npu==2.9.0.post2` | -| torchvision / torchaudio | `torchvision==0.24.0`，`torchaudio==2.9.0` | -| vLLM | 从 `vllm-project/vllm` 源码安装，默认分支 `v0.18.0` | -| vLLM Ascend | 从 `vllm-project/vllm-ascend` 源码安装，默认分支 `v0.18.0` | +| PyTorch | 默认 `torch==2.9.0`；可通过 `--torch_version` 配置 | +| torch-npu | 默认 `torch_npu==2.9.0.post2`；可通过 `--torch_npu_version` 配置 | +| torchvision / torchaudio | 默认 `torchvision==0.24.0`、`torchaudio==2.9.0`；覆盖 `--torch_version` 时必须同时显式传入两者 | +| vLLM | 从 `vllm-project/vllm` 源码安装，默认 `0.18.0`；可通过 `--vllm_version` 配置 | +| vLLM Ascend | 从 `vllm-project/vllm-ascend` 源码安装，默认 `0.18.0`；可通过 `--vllm_ascend_version` 配置 | | Megatron-LM | 源码 checkout，默认分支 `v0.15.3` | | MindSpeed | 源码 checkout，默认分支 `core_r0.15.3` | -| mcore-bridge | 来自 `modelscope/mcore-bridge` 的源码 checkout | +| mcore-bridge | PyPI 上的最新发布版 | | ms-swift | 来自 `modelscope/ms-swift` 的源码 checkout，默认分支 `main` | | ModelScope | 来自 `modelscope/modelscope` 的源码 checkout，默认分支 `master` | -| triton-ascend | CANN `8.5.*` 安装 `3.2.0`；CANN `9.0.0` 下载并本地安装 `3.2.1` wheel | +| triton-ascend | CANN `8.5.*` 默认 `3.2.0`；CANN `9.0.*` 默认 `3.2.1`；可通过 `--triton_ascend_version` 配置，并从 Triton Ascend PyPI 源安装 | ## 支持的 Tag 格式 通过 `docker/build_image.py --image_type ascend` 构建的镜像使用以下 tag 格式： ```text -${DOCKER_REGISTRY}:----- +${DOCKER_REGISTRY}:--torch_npu---- ``` | 字段 | 示例 | 说明 | | --- | --- | --- | | `swift-branch` | `main` | 构建镜像时使用的 ms-swift 分支 | -| `atlas-hardware` | `A2`、`A3`、`300I`、`A5` | 从 `--soc_version` 推导 | -| `python-tag` | `py311` | 从 `--python_version` 推导 | -| `cann-version-tag` | `CANN8.5.1`、`CANN9.0.0` | 从 CANN 基础镜像 tag 解析 | -| `os-tag` | `ubuntu22.04` | 从 CANN 基础镜像 tag 解析 | -| `arch` | `arm`、`x86` | 从宿主机架构或 `--arch` 推导 | +| `cann-version-tag` | `cann8.5.1`、`cann9.0.0` | 从 CANN 基础镜像 tag 解析 | +| `torch-npu-version` | `2.9.0.post2` | 来自 `--torch_npu_version`，默认 `2.9.0.post2` | +| `atlas-hardware` | `a2`、`a3`、`300i`、`a5` | 从 `--soc_version` 推导 | +| `os-tag` | `ubuntu22.04`、`openeuler24.03` | 从 CANN 基础镜像 tag 解析；避免不同 OS 的镜像 tag 冲突 | +| `python-tag` | `py3.11` | 从 CANN 基础镜像 tag 解析 | +| `arch` | `aarch64`、`x86_64` | 从宿主机架构或 `--arch` 推导 | ARM64 宿主机上的默认示例： ```text -${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm +${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64 ``` A2 / CANN 9.0.0 示例： ```text -${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm +${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 ``` ## 本地构建 @@ -85,6 +87,36 @@ python docker/build_image.py \ --soc_version ascend910b1 ``` +构建 openEuler 镜像。系统依赖层会自动使用 `yum`，Ubuntu 镜像继续使用 `apt-get`： + +```bash +python docker/build_image.py \ + --image_type ascend \ + --base_image quay.io/ascend/cann:8.5.1-a3-openeuler24.03-py3.11 \ + --soc_version ascend910_9391 +``` + +覆盖 PyTorch 版本组。`--torch_version` 必须与 `--torch_npu_version` 的基础版本一致；覆盖 PyTorch 时，必须显式传入匹配的 torchvision 和 torchaudio 版本： + +```bash +python docker/build_image.py \ + --image_type ascend \ + --torch_version 2.9.0 \ + --torch_npu_version 2.9.0.post2 \ + --torchvision_version 0.24.0 \ + --torchaudio_version 2.9.0 +``` + +覆盖 vLLM 版本组或 triton-ascend。vLLM 版本参数会选择对应 Git tag，例如 `0.18.0` 会选择 `v0.18.0`。 + +```bash +python docker/build_image.py \ + --image_type ascend \ + --vllm_version 0.18.0 \ + --vllm_ascend_version 0.18.0 \ + --triton_ascend_version 3.2.1 +``` + 需要时可以覆盖 Megatron 或 MindSpeed 源码分支： ```bash @@ -94,11 +126,11 @@ python docker/build_image.py \ --mindspeed_branch core_r0.15.3 ``` -如果构建时网络较慢，Linux 宿主机可以在根目录 `Dockerfile` 生成后使用 host network 构建： +如需手工构建生成后的根目录 `Dockerfile`，可使用： ```bash -docker build --network host \ - -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ +docker build \ + -t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \ -f Dockerfile . ``` @@ -119,7 +151,7 @@ docker run --rm -it \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /mnt/workspace:/mnt/workspace \ - ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \ bash ``` @@ -147,7 +179,8 @@ pip show ms-swift modelscope torch-npu triton-ascend ## 注意事项 - CANN、firmware 和 driver 版本必须互相兼容。 -- 这个 Dockerfile 对 CANN `8.5.*` 和 CANN `9.0.0` 使用不同的 `triton-ascend` 安装路径。 +- Ubuntu 基础镜像通过 `apt-get` 安装系统依赖；openEuler 基础镜像通过 `yum` 安装对应 RPM 包。 +- `triton-ascend` 从 `https://triton-ascend.osinfra.cn/pypi/simple` 安装；请选择与 CANN、Python 和架构兼容的版本。 - 该镜像面向 Ascend NPU 上的 ms-swift 工作流。依赖安装过程中引入且与 NPU runtime 冲突的 CUDA-only 包会被移除。 - 生产任务建议使用固定镜像 tag，不要依赖浮动分支名。 diff --git a/docker/build_image.py b/docker/build_image.py index 80e8ed06..18c1f2dd 100644 --- a/docker/build_image.py +++ b/docker/build_image.py @@ -469,20 +469,33 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy class AscendImageBuilder(StableGPUImageBuilder): + _DEFAULT_TORCH_VERSION = '2.9.0' + _DEFAULT_TORCHVISION_VERSION = '0.24.0' + _DEFAULT_TORCHAUDIO_VERSION = '2.9.0' + _DEFAULT_TORCH_NPU_VERSION = '2.9.0.post2' + _DEFAULT_VLLM_VERSION = '0.18.0' + _DEFAULT_VLLM_ASCEND_VERSION = '0.18.0' + _DEFAULT_TRITON_ASCEND_VERSIONS = { + '8.5': '3.2.0', + '9.0': '3.2.1', + } _CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$') _OS_TAG_PATTERN = re.compile(r'^[A-Za-z]+[0-9][0-9A-Za-z.]*$') + _PYTHON_TAG_PATTERN = re.compile(r'^py\d+\.\d+$', re.IGNORECASE) + _TORCH_NPU_VERSION_PATTERN = re.compile( + r'^(?P\d+\.\d+\.\d+)(?:\.post\d+)?$') @staticmethod def _normalize_arch(arch: str = None) -> str: arch = arch or platform.machine() arch = arch.lower() arch_mapping = { - 'x86': 'x86', - 'x86_64': 'x86', - 'amd64': 'x86', - 'arm': 'arm', - 'aarch64': 'arm', - 'arm64': 'arm', + 'x86': 'x86_64', + 'x86_64': 'x86_64', + 'amd64': 'x86_64', + 'arm': 'aarch64', + 'aarch64': 'aarch64', + 'arm64': 'aarch64', } if arch not in arch_mapping: raise ValueError(f'Unsupported architecture: {arch}. ' @@ -522,25 +535,107 @@ class AscendImageBuilder(StableGPUImageBuilder): cann_version = parts[0] os_tag = parts[2] + python_tag = parts[3] if not cls._CANN_VERSION_PATTERN.fullmatch(cann_version): raise ValueError(f'Invalid CANN version in Ascend base image tag: ' f'{cann_version}') if not cls._OS_TAG_PATTERN.fullmatch(os_tag): raise ValueError( f'Invalid OS tag in Ascend base image tag: {os_tag}') + if not cls._PYTHON_TAG_PATTERN.fullmatch(python_tag): + raise ValueError( + f'Invalid Python tag in Ascend base image tag: {python_tag}') - return cann_version, f'CANN{cann_version}', os_tag + return cann_version, f'CANN{cann_version}', os_tag, python_tag + + @staticmethod + def _get_os_family(os_tag: str) -> str: + os_tag = os_tag.lower() + if os_tag.startswith('ubuntu'): + return 'ubuntu' + if os_tag.startswith('openeuler'): + return 'openeuler' + raise ValueError(f'Unsupported Ascend base image OS tag: {os_tag}. ' + 'Supported OS families are Ubuntu and openEuler.') + + @classmethod + def _init_torch_versions(cls, args) -> None: + torch_version_specified = args.torch_version is not None + torchvision_version_specified = args.torchvision_version is not None + torchaudio_version_specified = args.torchaudio_version is not None + + if torch_version_specified: + if (not torchvision_version_specified + or not torchaudio_version_specified): + raise ValueError( + 'When overriding --torch_version for an Ascend image, also ' + 'pass matching --torchvision_version and ' + '--torchaudio_version.') + elif torchvision_version_specified or torchaudio_version_specified: + raise ValueError( + '--torchvision_version and --torchaudio_version require an ' + 'explicit --torch_version for an Ascend image.') + + args.torch_version = args.torch_version or cls._DEFAULT_TORCH_VERSION + args.torchvision_version = ( + args.torchvision_version or cls._DEFAULT_TORCHVISION_VERSION) + args.torchaudio_version = ( + args.torchaudio_version or cls._DEFAULT_TORCHAUDIO_VERSION) + args.torch_npu_version = ( + args.torch_npu_version or cls._DEFAULT_TORCH_NPU_VERSION) + + match = cls._TORCH_NPU_VERSION_PATTERN.fullmatch( + args.torch_npu_version) + if not match: + raise ValueError('Invalid --torch_npu_version. Expected ' + '.. or ' + '...post.') + if args.torch_version != match.group('torch_version'): + raise ValueError( + '--torch_version must exactly match the base version of ' + f'--torch_npu_version, got torch={args.torch_version} and ' + f'torch_npu={args.torch_npu_version}.') + + @classmethod + def _init_component_versions(cls, args) -> None: + args.vllm_version = args.vllm_version or cls._DEFAULT_VLLM_VERSION + args.vllm_ascend_version = ( + args.vllm_ascend_version or cls._DEFAULT_VLLM_ASCEND_VERSION) + args.vllm_git_ref = cls._get_vllm_git_ref(args.vllm_version) + args.vllm_ascend_git_ref = cls._get_vllm_git_ref( + args.vllm_ascend_version) + + if not args.triton_ascend_version: + cann_series = '.'.join(args.cann_version.split('.')[:2]) + try: + args.triton_ascend_version = ( + cls._DEFAULT_TRITON_ASCEND_VERSIONS[cann_series]) + except KeyError as e: + raise ValueError('No default triton-ascend version for CANN ' + f'{args.cann_version}. Please pass ' + '--triton_ascend_version explicitly.') from e + + @staticmethod + def _get_vllm_git_ref(version: str) -> str: + return version if version.startswith('v') else f'v{version}' def init_args(self, args) -> Any: if not args.base_image: # Reuse the prebuilt vllm-ascend image to avoid rebuilding its stack. args.base_image = 'quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11' + self._init_torch_versions(args) args.arch = self._normalize_arch(args.arch) args.atlas_hardware = self._get_atlas_hardware(args.soc_version) - args.cann_version, args.cann_version_tag, args.os_tag = ( - self._get_cann_os_tags(args.base_image)) + (args.cann_version, args.cann_version_tag, args.os_tag, + args.ascend_python_tag) = ( + self._get_cann_os_tags(args.base_image)) + self._get_os_family(args.os_tag) + self._init_component_versions(args) return super().init_args(args) + def _generate_python_tag(self, _python_version: str) -> str: + return self.args.ascend_python_tag + def generate_dockerfile(self) -> str: extra_content = """ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy @@ -550,6 +645,19 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy content = content.replace('{base_image}', self.args.base_image) content = content.replace('{soc_version}', self.args.soc_version) content = content.replace('{cann_version}', self.args.cann_version) + content = content.replace('{torch_version}', + self.args.torch_version) + content = content.replace('{torchvision_version}', + self.args.torchvision_version) + content = content.replace('{torchaudio_version}', + self.args.torchaudio_version) + content = content.replace('{torch_npu_version}', + self.args.torch_npu_version) + content = content.replace('{vllm_git_ref}', self.args.vllm_git_ref) + content = content.replace('{vllm_ascend_git_ref}', + self.args.vllm_ascend_git_ref) + content = content.replace('{triton_ascend_version}', + self.args.triton_ascend_version) content = content.replace('{extra_content}', extra_content) content = content.replace('{cur_time}', formatted_time) content = content.replace('{install_ms_deps}', 'False') @@ -563,11 +671,12 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy return content def image(self) -> str: - return ( - f'{docker_registry}:{self.args.swift_branch}-' - f'{self.args.atlas_hardware}-{self.args.python_tag}-' - f'{self.args.cann_version_tag}-{self.args.os_tag}-{self.args.arch}' - ) + tag = (f'{self.args.swift_branch}-{self.args.cann_version_tag}-' + f'torch_npu{self.args.torch_npu_version}-' + f'{self.args.atlas_hardware}-{self.args.os_tag}-' + f'{self.args.python_tag}-' + f'{self.args.arch}') + return f'{docker_registry}:{tag.lower()}' def push(self): return 0 @@ -579,12 +688,15 @@ parser.add_argument('--image_type', type=str) parser.add_argument('--python_version', type=str, default='3.12.13') parser.add_argument('--ubuntu_version', type=str, default='22.04') parser.add_argument('--torch_version', type=str, default=None) +parser.add_argument('--torch_npu_version', type=str, default=None) parser.add_argument('--torchvision_version', type=str, default=None) parser.add_argument('--cuda_version', type=str, default=None) parser.add_argument('--ci_image', type=int, default=0) parser.add_argument('--torchaudio_version', type=str, default=None) parser.add_argument('--tf_version', type=str, default=None) parser.add_argument('--vllm_version', type=str, default=None) +parser.add_argument('--vllm_ascend_version', type=str, default=None) +parser.add_argument('--triton_ascend_version', type=str, default=None) parser.add_argument('--lmdeploy_version', type=str, default=None) parser.add_argument('--flashattn_version', type=str, default=None) parser.add_argument('--autogptq_version', type=str, default=None) diff --git a/tests/msdatasets/test_stream_load.py b/tests/msdatasets/test_stream_load.py index 87bbf096..4ed0c3f3 100644 --- a/tests/msdatasets/test_stream_load.py +++ b/tests/msdatasets/test_stream_load.py @@ -44,15 +44,12 @@ class TestStreamLoad(unittest.TestCase): pass self.assertIs(HfFileSystem._open, hf_datasets_util._hf_fs_open) - self.assertIsNot( - hf_datasets_util._hf_fs_open_original, - hf_datasets_util._hf_fs_open) - self.assertIs( - HfFileSystem.__init__, - hf_datasets_util._hf_fs_init_with_cookie) - self.assertIsNot( - hf_datasets_util._hf_fs_init_original, - hf_datasets_util._hf_fs_init_with_cookie) + self.assertIsNot(hf_datasets_util._hf_fs_open_original, + hf_datasets_util._hf_fs_open) + self.assertIs(HfFileSystem.__init__, + hf_datasets_util._hf_fs_init_with_cookie) + self.assertIsNot(hf_datasets_util._hf_fs_init_original, + hf_datasets_util._hf_fs_init_with_cookie) finally: HfFileSystem._open = hf_fs_open_before HfFileSystem.__init__ = hf_fs_init_before @@ -75,7 +72,8 @@ class TestStreamLoad(unittest.TestCase): 'load_dataset', side_effect=RuntimeError('load failed')): with self.assertRaises(RuntimeError): - with hf_datasets_util.load_dataset_with_ctx(streaming=True): + with hf_datasets_util.load_dataset_with_ctx( + streaming=True): pass self.assertIs(HfFileSystem._open, hf_fs_open_clean)