update npu dockerfile (#1755)

This commit is contained in:
addsubmuldiv
2026-07-23 16:30:26 +08:00
committed by GitHub
parent 61571c0169
commit 2905da1220
5 changed files with 288 additions and 110 deletions

View File

@@ -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 && \

View File

@@ -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-<arch>`
- 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-<arch>`
- 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}:<swift-branch>-<atlas-hardware>-<python-tag>-<cann-version-tag>-<os-tag>-<arch>
${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<torch-npu-version>-<atlas-hardware>-<os-tag>-<python-tag>-<arch>
```
| 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.

View File

@@ -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-<arch>`
- 支持的基础 OSUbuntu 和 openEuler由 CANN 基础镜像 tag 选择
- 默认输出 tag`${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-<arch>`
- 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}:<swift-branch>-<atlas-hardware>-<python-tag>-<cann-version-tag>-<os-tag>-<arch>
${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<torch-npu-version>-<atlas-hardware>-<os-tag>-<python-tag>-<arch>
```
| 字段 | 示例 | 说明 |
| --- | --- | --- |
| `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不要依赖浮动分支名。

View File

@@ -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<torch_version>\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 '
'<major>.<minor>.<patch> or '
'<major>.<minor>.<patch>.post<revision>.')
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)

View File

@@ -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)