Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-08-29 10:08:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

Update Ascend image build and refresh runtime stack (#1781)

Browse Source
This commit is contained in:
addsubmuldiv
2026-08-24 12:11:31 +08:00
committed by GitHub
parent 34dadf9d34
commit 772955e670
8 changed files with 673 additions and 590 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
docker/Dockerfile.ascend
View File

@@ -1,176 +0,0 @@
FROM {base_image}
ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_DEFAULT_TIMEOUT=300 \
PIP_RETRIES=10 \
SOC_VERSION={soc_version} \
CANN_VERSION={cann_version}
# Build-time only (ARG does not persist into the image).
ARG PIP_EXTRA_INDEX_URL=https://pypi.org/simple
SHELL ["/bin/bash", "-c"]
# ---------- System dependencies ----------
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 install.trusted-host mirrors.aliyun.com
{extra_content}
# ---------- 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 {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) && \
export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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=={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
cd vllm-ascend && pip install -v -e . && cd ..
# ---------- 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
# ---------- Install training-side repositories ----------
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 . && \
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 ----------
# x86: must force-install the CPU build from pytorch.org/whl/cpu
# aarch64: PyPI only provides the CPU build, so install it directly from the Aliyun mirror
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "x86_64" ]; then \
pip install --no-cache-dir --force-reinstall --no-deps \
--index-url https://download.pytorch.org/whl/cpu \
torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
else \
pip install --no-cache-dir --force-reinstall --no-deps \
torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
fi && \
pip install --no-cache-dir --force-reinstall --no-deps \
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) ----------
RUN pip uninstall -y flashinfer tvm-ffi torch-c-dlpack-ext 2>/dev/null || true
ARG INSTALL_MS_DEPS={install_ms_deps}
ENV MEGATRON_LM_PATH=/Megatron-LM
ENV PYTHONPATH=/Megatron-LM:${PYTHONPATH}
# install dependencies
COPY requirements /var/modelscope
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
pip uninstall ms-swift modelscope -y && pip install --no-cache-dir pip==23.* -U && \
if [ "$INSTALL_MS_DEPS" = "True" ]; then \
pip install --no-cache-dir omegaconf==2.0.6 && \
pip install 'editdistance==0.8.1' && \
pip install --no-cache-dir 'cython>=3.0' versioneer 'numpy>=2.0' && \
pip install --no-cache-dir -r /var/modelscope/framework.txt && \
pip install --no-cache-dir -r /var/modelscope/audio.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
pip install --no-cache-dir -r /var/modelscope/tests.txt && \
pip install --no-cache-dir -r /var/modelscope/server.txt && \
pip install --no-cache-dir https://modelscope.oss-cn-beijing.aliyuncs.com/packages/imageio_ffmpeg-0.4.9-py3-none-any.whl --no-dependencies --force && \
pip install --no-cache-dir 'scipy' && \
pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn && \
pip install --no-cache-dir 'decord>=0.6.0' mpi4py paint_ldm ipykernel fasttext && \
pip install --no-cache-dir ipywidgets && \
pip install --no-cache-dir 'blobfile>=1.0.5' && \
pip uninstall MinDAEC -y && \
pip install https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/MinDAEC-0.0.2-py3-none-any.whl && \
pip cache purge; \
else \
pip install --no-cache-dir -r /var/modelscope/framework.txt && \
pip cache purge; \
fi
ARG CUR_TIME={cur_time}
RUN echo $CUR_TIME
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
pip install --no-cache-dir --no-build-isolation OpenCC
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 && \
pip install --no-cache-dir -U funasr scikit-learn && \
pip install --no-cache-dir -U qwen_vl_utils qwen_omni_utils librosa 'timm>=0.9.0' transformers accelerate peft trl safetensors && \
cd /ms-swift && pip install --no-cache-dir -e '.[llm]' && \
pip install --no-cache-dir -e '.[eval]' && pip install evalscope -U --no-dependencies && \
pip install --no-cache-dir "numpy>=2.0" --force-reinstall --no-deps && \
pip install ms-agent -U --no-dependencies && \
cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b {modelscope_branch} --single-branch https://github.com/modelscope/modelscope.git && \
cd modelscope && pip install . -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
cd / && rm -fr /tmp/modelscope && \
pip install --no-cache-dir diffusers && \
pip install --no-cache-dir omegaconf==2.3.0 && \
pip cache purge
# ---------- 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; \
export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}"; \
pip uninstall -y triton || true; \
pip uninstall -y triton-ascend || true; \
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 && \
echo 'set +H' >> /root/.bashrc
ENV SETUPTOOLS_USE_DISTUTILS=stdlib
ENV VLLM_USE_MODELSCOPE=True
ENV LMDEPLOY_USE_MODELSCOPE=True
ENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope/hub
# Show install results
RUN pip list
WORKDIR /workspace
CMD ["/bin/bash"]

192
docker/OVERVIEW.ascend.md
View File

@@ -1,192 +0,0 @@
# ms-swift Ascend
> English | [中文](./OVERVIEW.ascend.zh.md)
ms-swift Ascend images provide a ready-to-use ms-swift environment for Huawei Ascend Atlas NPUs. The images are built on top of the Ascend CANN container images and include the Python, CANN, PyTorch NPU, vLLM Ascend, Megatron, MindSpeed, mcore-bridge, ms-swift, and ModelScope runtime components needed for Ascend inference and training workflows.
## Quick Reference
- Base image: `quay.io/ascend/cann:<cann-version>-<hardware>-<os>-py<python-version>`
- 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`
- 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`
## Image Contents
The Ascend Dockerfile installs and configures:
| Component | Version / Source |
| --- | --- |
| 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` 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 | 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 | 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>-<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 |
| `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-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64
```
A2 / CANN 9.0.0 example:
```text
${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64
```
## Build Locally
Set the target registry first. The build script renders `docker/Dockerfile.ascend` into the root `Dockerfile`, builds it, and skips push for Ascend images.
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend
```
Build an A2 / CANN 9.0.0 image:
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend \
--base_image quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11 \
--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
python docker/build_image.py \
--image_type ascend \
--megatron_branch v0.15.3 \
--mindspeed_branch core_r0.15.3
```
To run the rendered Dockerfile manually, use:
```bash
docker build \
-t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
-f Dockerfile .
```
## Run An Ascend Container
The host must have a compatible Ascend driver and firmware installed. The container uses the host NPU devices and driver libraries.
```bash
docker run --rm -it \
--name ms_swift_ascend \
--device /dev/davinci0 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \
-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-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
bash
```
Inside the container, verify the NPU and Python packages:
```bash
npu-smi info
python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)"
python -c "import vllm, vllm_ascend; print('vllm ascend ok')"
pip show ms-swift modelscope torch-npu triton-ascend
```
## Environment Variables
| Variable | Value |
| --- | --- |
| `SOC_VERSION` | Selected Ascend SoC version, for example `ascend910b1` or `ascend910_9391` |
| `CANN_VERSION` | Parsed from the base image tag |
| `MEGATRON_LM_PATH` | `/Megatron-LM` |
| `PYTHONPATH` | includes `/Megatron-LM` |
| `VLLM_USE_MODELSCOPE` | `True` |
| `LMDEPLOY_USE_MODELSCOPE` | `True` |
| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` |
## Notes
- CANN, firmware, and driver versions must be compatible with each other.
- 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.
## License
ms-swift and ModelScope components follow their upstream repository licenses. CANN, MindSpeed, torch-npu, vLLM Ascend, and other pre-installed third-party components are subject to their own upstream licenses.

189
docker/OVERVIEW.ascend.zh.md
View File

@@ -1,189 +0,0 @@
# ms-swift Ascend
> [English](./OVERVIEW.ascend.md) | 中文
ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms-swift 运行环境。镜像基于 Ascend CANN 容器镜像构建，包含 Ascend 推理和训练工作流所需的 Python、CANN、PyTorch NPU、vLLM Ascend、Megatron、MindSpeed、mcore-bridge、ms-swift 以及 ModelScope 运行组件。
## 快速参考
- 基础镜像：`quay.io/ascend/cann:<cann-version>-<hardware>-<os>-py<python-version>`
- 构建模板：`docker/Dockerfile.ascend`
- 构建入口：`docker/build_image.py --image_type ascend`
- 默认基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`
- 支持的基础 OS：Ubuntu 和 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`
## 镜像内容
Ascend Dockerfile 会安装和配置：
| 组件 | 版本 / 来源 |
| --- | --- |
| CANN | 继承自选定的 `quay.io/ascend/cann` 基础镜像 |
| Python | 继承自基础镜像 tag，例如 `py3.11` |
| 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 | PyPI 上的最新发布版 |
| ms-swift | 来自 `modelscope/ms-swift` 的源码 checkout，默认分支 `main` |
| ModelScope | 来自 `modelscope/modelscope` 的源码 checkout，默认分支 `master` |
| 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>-<cann-version-tag>-torch_npu<torch-npu-version>-<atlas-hardware>-<os-tag>-<python-tag>-<arch>
```
| 字段 | 示例 | 说明 |
| --- | --- | --- |
| `swift-branch` | `main` | 构建镜像时使用的 ms-swift 分支 |
| `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-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64
```
A2 / CANN 9.0.0 示例：
```text
${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64
```
## 本地构建
先设置目标镜像仓库。构建脚本会把 `docker/Dockerfile.ascend` 渲染成根目录 `Dockerfile`，然后执行构建；Ascend 镜像分支当前不执行 push。
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend
```
构建 A2 / CANN 9.0.0 镜像：
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend \
--base_image quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11 \
--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
python docker/build_image.py \
--image_type ascend \
--megatron_branch v0.15.3 \
--mindspeed_branch core_r0.15.3
```
如需手工构建生成后的根目录 `Dockerfile`，可使用：
```bash
docker build \
-t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
-f Dockerfile .
```
## 运行 Ascend 容器
宿主机需要提前安装兼容的 Ascend driver 和 firmware。容器通过挂载宿主机 NPU 设备和 driver 库使用昇腾硬件。
```bash
docker run --rm -it \
--name ms_swift_ascend \
--device /dev/davinci0 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \
-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-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
bash
```
进入容器后可以验证 NPU 和 Python 包：
```bash
npu-smi info
python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)"
python -c "import vllm, vllm_ascend; print('vllm ascend ok')"
pip show ms-swift modelscope torch-npu triton-ascend
```
## 环境变量
| 变量 | 值 |
| --- | --- |
| `SOC_VERSION` | 选定的 Ascend SoC，例如 `ascend910b1` 或 `ascend910_9391` |
| `CANN_VERSION` | 从基础镜像 tag 解析得到 |
| `MEGATRON_LM_PATH` | `/Megatron-LM` |
| `PYTHONPATH` | 包含 `/Megatron-LM` |
| `VLLM_USE_MODELSCOPE` | `True` |
| `LMDEPLOY_USE_MODELSCOPE` | `True` |
| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` |
## 注意事项
- CANN、firmware 和 driver 版本必须互相兼容。
- 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，不要依赖浮动分支名。
## License
ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MindSpeed、torch-npu、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。

188
docker/ascend/Dockerfile.ascend Normal file
View File

@@ -0,0 +1,188 @@
FROM {base_image}
ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_DEFAULT_TIMEOUT=300 \
PIP_RETRIES=10 \
SOC_VERSION={soc_version} \
CANN_VERSION={cann_version}
# Build-time only (ARG does not persist into the image).
ARG PIP_EXTRA_INDEX_URL={pip_extra_index_url}
ARG PYPI_OFFICIAL_INDEX_URL={pypi_official_index_url}
SHELL ["/bin/bash", "-c"]
# ---------- System dependencies ----------
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://repo.huaweicloud.com/repository/pypi/simple && \
pip config set install.trusted-host repo.huaweicloud.com
# ---------- Install PyTorch and TorchNPU ----------
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "x86_64" ]; then \
TORCH_INDEX=https://download.pytorch.org/whl/cpu; \
else \
TORCH_INDEX=https://repo.huaweicloud.com/repository/pypi/simple; \
fi && \
PIP_EXTRA_INDEX_URL= pip install --no-cache-dir --index-url "$TORCH_INDEX" \
torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version} && \
PIP_EXTRA_INDEX_URL= pip install --no-cache-dir \
--index-url "${PIP_EXTRA_INDEX_URL}" \
--extra-index-url "${PYPI_OFFICIAL_INDEX_URL}" \
torch_npu=={torch_npu_version}
# ---------- Install triton-ascend before vLLM ----------
RUN set -eux; \
ASCEND_PYPI_INDEX_URL="${PIP_EXTRA_INDEX_URL}"; \
pip uninstall -y triton || true; \
pip uninstall -y triton-ascend || true; \
PIP_INDEX_URL="${ASCEND_PYPI_INDEX_URL}" \
PIP_EXTRA_INDEX_URL="${PYPI_OFFICIAL_INDEX_URL}" \
pip install --no-cache-dir --force-reinstall \
triton-ascend=={triton_ascend_version}
# ---------- Install vLLM from the official source tag ----------
# Keep both checkouts: editable installs resolve their Python packages from these paths at runtime.
RUN git clone --depth 1 --branch {vllm_git_ref} \
https://github.com/vllm-project/vllm /vllm && \
cd /vllm && \
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 && \
VLLM_TARGET_DEVICE=empty pip install --no-cache-dir -v -e .
# ---------- Install vLLM Ascend from the official source tag ----------
RUN git clone --depth 1 --branch {vllm_ascend_git_ref} \
https://github.com/vllm-project/vllm-ascend /vllm-ascend && \
cd /vllm-ascend && \
git submodule update --init --recursive && \
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 && \
ASCEND_PYPI_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
PIP_INDEX_URL="${ASCEND_PYPI_INDEX_URL}" \
PIP_EXTRA_INDEX_URL="${PYPI_OFFICIAL_INDEX_URL}" \
pip install --no-cache-dir -v -e .
# ---------- 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
# ---------- Install core training components ----------
# Install the published ModelScope package; do not bind the image to source-repository snapshots.
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 . && \
pip install --no-cache-dir mcore-bridge -i https://pypi.org/simple/ -U && \
pip install --no-cache-dir -U modelscope && \
cd /ms-swift && pip install --no-cache-dir -e .
# ---------- Pin torch to the correct version + torch_npu ----------
# x86: must force-install the CPU build from pytorch.org/whl/cpu
# aarch64: PyPI only provides the CPU build, so install it directly from the Huawei Cloud mirror
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "x86_64" ]; then \
pip install --no-cache-dir --force-reinstall --no-deps \
--index-url https://download.pytorch.org/whl/cpu \
torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
else \
pip install --no-cache-dir --force-reinstall --no-deps \
torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
fi && \
PIP_EXTRA_INDEX_URL= pip install --no-cache-dir --force-reinstall --no-deps \
--index-url "${PIP_EXTRA_INDEX_URL}" \
--extra-index-url "${PYPI_OFFICIAL_INDEX_URL}" \
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) ----------
RUN pip uninstall -y flashinfer tvm-ffi torch-c-dlpack-ext 2>/dev/null || true
ENV MEGATRON_LM_PATH=/Megatron-LM
ENV PYTHONPATH=/Megatron-LM:${PYTHONPATH}
ARG CUR_TIME={cur_time}
RUN echo $CUR_TIME
# ---------- Install Swift extras and auxiliary tools ----------
RUN export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL}" && \
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 && \
pip install --no-cache-dir -U pip && \
pip install --no-cache-dir -U icecream soundfile pybind11 py-spy funasr scikit-learn && \
pip install --no-cache-dir -U qwen_vl_utils qwen_omni_utils 'librosa<1.0' 'timm>=0.9.0' transformers accelerate peft trl safetensors && \
pip install --no-cache-dir -r /ms-swift/requirements/eval.txt && \
pip install --no-cache-dir -U evalscope --no-dependencies && \
pip install --no-cache-dir 'numpy>=2.0' --force-reinstall --no-deps && \
pip install --no-cache-dir -U ms-agent --no-dependencies && \
pip install --no-cache-dir --no-build-isolation OpenCC && \
pip install --no-cache-dir diffusers omegaconf==2.3.0 && \
pip cache purge
# ---------- 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 -U "deepspeed>=0.19" ray liger_kernel pre-commit && \
pip cache purge
# ---------- Install flash-linear-attention (FLA) ----------
# Keep the checkout because the editable install points to this source tree.
RUN git clone --depth 1 --branch {fla_version} https://github.com/fla-org/flash-linear-attention.git /flash-linear-attention && \
cd /flash-linear-attention && \
pip install --no-cache-dir -e .
# ---------- Re-pin triton-ascend after all editable installs ----------
# vllm-ascend may otherwise replace the CANN-specific version selected above.
RUN set -eux; \
ASCEND_PYPI_INDEX_URL="${PIP_EXTRA_INDEX_URL}"; \
pip uninstall -y triton || true; \
pip uninstall -y triton-ascend || true; \
PIP_INDEX_URL="${ASCEND_PYPI_INDEX_URL}" \
PIP_EXTRA_INDEX_URL="${PYPI_OFFICIAL_INDEX_URL}" \
pip install --no-cache-dir --force-reinstall \
triton-ascend=={triton_ascend_version}
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 && \
echo 'set +H' >> /root/.bashrc
ENV VLLM_USE_MODELSCOPE=True
ENV LMDEPLOY_USE_MODELSCOPE=True
ENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope/hub
# Show install results and verify that editable packages resolve from retained source trees.
RUN pip list && \
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -c "import importlib.metadata as m; import importlib.util as u; import os; expected = {'vllm': '/vllm', 'vllm_ascend': '/vllm-ascend', 'swift': '/ms-swift', 'megatron': '/Megatron-LM', 'mindspeed': '/MindSpeed', 'fla': '/flash-linear-attention'}; specs = {name: u.find_spec(name) for name in expected}; origins = {name: (spec.origin or next(iter(spec.submodule_search_locations or []), None) if spec else None) for name, spec in specs.items()}; assert all(path and os.path.commonpath([os.path.realpath(path), root]) == root for name, root in expected.items() for path in [origins[name]]), origins; assert u.find_spec('modelscope') and m.version('modelscope'); assert m.version('triton-ascend') == '{triton_ascend_version}', m.version('triton-ascend'); print(origins)"
WORKDIR /workspace
CMD ["/bin/bash"]

203
docker/ascend/OVERVIEW.ascend.md Normal file
View File

@@ -0,0 +1,203 @@
# ms-swift Ascend
> English | [中文](./OVERVIEW.ascend.zh.md)
ms-swift Ascend images provide a ready-to-use ms-swift environment for Huawei Ascend Atlas NPUs. The images are built on top of the Ascend CANN container images and include the Python, CANN, TorchNPU, vLLM Ascend, FLA, Megatron, MindSpeed, mcore-bridge, ms-swift, and ModelScope runtime components needed for Ascend inference and training workflows.
## Quick Reference
- Base image: `quay.io/ascend/cann:<cann-version>-<hardware>-<os>-py<python-version>`
- Build template: `docker/ascend/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`
- 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`
## Image Contents
The Ascend Dockerfile installs and configures:
```text
| Component | Version / Source |
|--------------------------|------------------|
| CANN | inherited from the selected `quay.io/ascend/cann` base image; configurable with `--base_image` |
| Python | inherited from the base image tag, for example `py3.11`; configure it by selecting a matching tag with `--base_image` |
| PyTorch | `torch==2.9.0` by default; configurable with `--torch_version` (also pass matching `--torchvision_version` and `--torchaudio_version`) |
| TorchNPU | `torch_npu==2.9.0.post2` by default; configurable with `--torch_npu_version`, which must match the PyTorch base version |
| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` by default; configurable with `--torchvision_version` and `--torchaudio_version` when overriding `--torch_version` |
| vLLM | Official `vllm-project/vllm` source tag `v0.18.0` by default; configurable with `--vllm_version`; built with the empty target device |
| vLLM Ascend | Official `vllm-project/vllm-ascend` source tag `v0.18.0` by default; configurable with `--vllm_ascend_version`; initializes submodules; resolves build and runtime dependencies from Huawei Cloud Ascend PyPI with official PyPI as fallback |
| FLA | source checkout from `fla-org/flash-linear-attention`, default branch `main`; configurable with `--fla_version` using a branch or release tag |
| Megatron-LM | source checkout, default branch `v0.15.3`; configurable with `--megatron_branch` |
| MindSpeed | source checkout, default branch `core_r0.15.3`; configurable with `--mindspeed_branch` |
| mcore-bridge | latest release from PyPI; no dedicated version option in the Ascend build entrypoint |
| ms-swift | source checkout from `modelscope/ms-swift`, default branch `main`; configurable with `--swift_branch` |
| DeepSpeed | newest package satisfying `deepspeed>=0.19`; `TORCH_DEVICE_BACKEND_AUTOLOAD=0` is scoped to its build-time install command |
| ModelScope | latest published package resolved by `pip install -U modelscope`; the Ascend image does not clone ModelScope or modelscope-hub source repositories |
| triton-ascend | Installed before vLLM and re-pinned after all editable installs, using Huawei Cloud Ascend PyPI as the primary index and official PyPI as fallback; CANN `8.5.*` defaults to `3.2.0`, CANN `9.0.*` to `3.2.1`; configurable with `--triton_ascend_version` |
```
## Image Tag Description
Images built by `docker/build_image.py --image_type ascend` use this tag format:
The published-tag index is maintained in [`docker/ascend/supported_tags.md`](./supported_tags.md).
```text
${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<TorchNPU-version>-<hardware-tag>-<os-tag>-<python-tag>-<arch>
```
```text
| Field | Example | Description |
|--------------------|---------------------------------|-------------|
| `swift-branch` | `main` | ms-swift branch used during image build |
| `cann-version-tag` | `cann8.5.1`, `cann9.0.0` | Parsed from the CANN base image tag |
| `TorchNPU-version` | `2.9.0.post2` | From `--torch_npu_version`; defaults to `2.9.0.post2` |
| `hardware-tag` | `910b`, `a3`, `950` | Parsed directly from the CANN base-image tag. The current Atlas platform names are `a2`, `a3`, and `a5`; these are documentation labels only. |
| `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` |
```
A2 / CANN 9.0.0 example (`910b` CANN hardware tag):
```text
${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-910b-ubuntu22.04-py3.11-aarch64
```
## Latest Images
The latest A2 and A3 images are hosted at [quay.io/ascend/ms-swift](https://quay.io/repository/ascend/ms-swift?tab=tags).
**Device / CANN Base Image / OS / Image Tag / Dockerfile**
- A3 — 9.1.0 — openEuler 24.03 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-openeuler24.03-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A3 — 9.1.0 — Ubuntu 22.04 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-ubuntu22.04-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A2 — 9.1.0 — openEuler 24.03 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-openeuler24.03-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A2 — 9.1.0 — Ubuntu 22.04 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-ubuntu22.04-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
## Build Locally
Set the target registry first. The build script renders `docker/ascend/Dockerfile.ascend` into the root `Dockerfile`, builds it, and skips push for Ascend images.
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend
```
Complete version-pinned reference (CANN 9.1.0, Atlas A2, Ubuntu 22.04,
Python 3.12, ARM64):
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend \
--base_image quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12 \
--soc_version ascend910b1 \
--arch arm \
--torch_version 2.10.0 \
--torch_npu_version 2.10.0.post2 \
--torchvision_version 0.25.0 \
--torchaudio_version 2.10.0 \
--vllm_version v0.23.0 \
--vllm_ascend_version v0.23.0 \
--fla_version main \
--triton_ascend_version 3.2.2 \
--modelscope_branch master \
--swift_branch v4.5.2 \
--megatron_branch core_v0.16.0 \
--mindspeed_branch core_r0.16.0
```
## Custom Build Parameters
Use `--image_type ascend` to select the Ascend builder. The following options
are consumed by the current Ascend build path:
- `--base_image` (default: `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`): selects CANN, hardware, OS, and Python. The tag must follow `<cann-version>-<hardware>-<os>-py<python-version>`.
- `--soc_version` (default: `ascend910_9391`): sets the target SoC for the vLLM Ascend build and the runtime `SOC_VERSION`; it must match the target hardware. The [vLLM Ascend v0.23.0 installation guide](https://docs.vllm.ai/projects/ascend/en/v0.23.0/installation.html) provides the reference values below. These are upstream vLLM Ascend build targets, not a claim that every hardware path has been verified by this Dockerfile:
- Atlas A2: `ascend910b1`
- Atlas A3: `ascend910_9391`
- Atlas 300I DUO: `ascend310p1`
- Atlas 950DT: `ascend950dt_9582`
- `--arch` (default: detected from the host): accepts `arm` or `x86`; normalized to `aarch64` or `x86_64` in the output tag.
- `--torch_version` (default: `2.9.0`): selects PyTorch. If overridden, matching `--torchvision_version` and `--torchaudio_version` are required.
- `--torch_npu_version` (default: `2.9.0.post2`): its base version must exactly match `--torch_version`.
- `--torchvision_version` (default: `0.24.0`): selects torchvision; required when `--torch_version` is overridden.
- `--torchaudio_version` (default: `2.9.0`): selects torchaudio; required when `--torch_version` is overridden.
- `--vllm_version` (default: `0.18.0`): selects the official vLLM source tag; values with or without the leading `v` are accepted.
- `--vllm_ascend_version` (default: `0.18.0`): selects the official vLLM Ascend source tag; values with or without the leading `v` are accepted.
- `--fla_version` (default: `main`): selects the FLA branch or release tag.
- `--triton_ascend_version` (CANN-specific default): uses `3.2.0` for CANN 8.5 and `3.2.1` for CANN 9.0. Other CANN series, including 9.1, require an explicit value.
- `--pip_extra_index_url` (default: Huawei Cloud Ascend PyPI): sets the primary package index for Ascend-specific dependencies.
- `--pypi_official_index_url` (default: `https://pypi.org/simple`): sets the fallback package index.
- `--swift_branch` (default: `main`): selects the ms-swift source branch or tag and is included in the output image tag.
- `--megatron_branch` (default: `v0.15.3`): selects the Megatron-LM source branch or tag.
- `--mindspeed_branch` (default: `core_r0.15.3`): selects the MindSpeed source branch or tag.
`--python_version` does not override the Python version for Ascend images;
select it through `--base_image`. `--modelscope_branch` is accepted by the
shared argument parser but is not consumed by the Ascend Dockerfile, which
installs the latest published ModelScope package with `pip install -U
modelscope`.
## Run An Ascend Container
The host must have a compatible Ascend driver and firmware installed. The container uses the host NPU devices and driver libraries.
```bash
docker run --rm -it \
--name ms_swift_ascend \
--device /dev/davinci0 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \
-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-cann9.0.0-torch_npu2.9.0.post2-910b-ubuntu22.04-py3.11-aarch64 \
bash
```
Inside the container, verify the NPU and Python packages:
```bash
npu-smi info
python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)"
python -c "import vllm, vllm_ascend; print('vllm ascend ok')"
pip show ms-swift modelscope mcore-bridge torch-npu triton-ascend
```
## Environment Variables
```text
| Variable | Value |
|---------------------------|-------|
| `SOC_VERSION` | Selected Ascend SoC version, for example `ascend910b1` or `ascend910_9391`; configurable with `--soc_version`; used as the vLLM Ascend build target and retained in the runtime environment (the tag hardware comes from `--base_image`) |
| `CANN_VERSION` | Parsed from the base image tag |
| `MEGATRON_LM_PATH` | `/Megatron-LM` |
| `PYTHONPATH` | includes `/Megatron-LM` |
| `VLLM_USE_MODELSCOPE` | `True` |
| `LMDEPLOY_USE_MODELSCOPE` | `True` |
| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` |
```
## Notes
- CANN, firmware, and driver versions must be compatible with each other.
- Ubuntu base images install system dependencies through `apt-get`; openEuler base images install the corresponding RPM packages through `yum`.
- `triton-ascend` and the dependencies of the editable vLLM Ascend install use `https://mirrors.huaweicloud.com/ascend/repos/pypi` as the primary index and `https://pypi.org/simple` as fallback. pip evaluates candidates from all configured indexes rather than enforcing strict repository priority. Select versions 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.
## License
ms-swift and ModelScope components follow their upstream repository licenses. CANN, MindSpeed, TorchNPU, vLLM Ascend, and other pre-installed third-party components are subject to their own upstream licenses.

200
docker/ascend/OVERVIEW.ascend.zh.md Normal file
View File

@@ -0,0 +1,200 @@
# ms-swift Ascend
> [English](./OVERVIEW.ascend.md) | 中文
ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms-swift 运行环境。镜像基于 Ascend CANN 容器镜像构建，包含 Ascend 推理和训练工作流所需的 Python、CANN、TorchNPU、vLLM Ascend、FLA、Megatron、MindSpeed、mcore-bridge、ms-swift 以及 ModelScope 运行组件。
## 快速参考
- 基础镜像：`quay.io/ascend/cann:<cann-version>-<hardware>-<os>-py<python-version>`
- 构建模板：`docker/ascend/Dockerfile.ascend`
- 构建入口：`docker/build_image.py --image_type ascend`
- 默认基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`
- 支持的基础 OS：Ubuntu 和 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`
## 镜像内容
Ascend Dockerfile 会安装和配置：
```text
| 组件 | 版本 / 来源 |
|----------------------------|------------|
| CANN | 继承自选定的 `quay.io/ascend/cann` 基础镜像；可通过 `--base_image` 配置 |
| Python | 继承自基础镜像 tag，例如 `py3.11`；通过 `--base_image` 选择匹配的 Python tag |
| PyTorch | 默认 `torch==2.9.0`；可通过 `--torch_version` 配置（同时需要传入匹配的 `--torchvision_version` 和 `--torchaudio_version`） |
| TorchNPU | 默认 `torch_npu==2.9.0.post2`；可通过 `--torch_npu_version` 配置，但必须与 PyTorch 基础版本匹配 |
| torchvision / torchaudio | 默认 `torchvision==0.24.0`、`torchaudio==2.9.0`；覆盖 `--torch_version` 时可通过 `--torchvision_version` 和 `--torchaudio_version` 配置 |
| vLLM | 默认从官方 `vllm-project/vllm` 源码 tag `v0.18.0` 构建；可通过 `--vllm_version` 配置；使用 empty target device |
| vLLM Ascend | 默认从官方 `vllm-project/vllm-ascend` 源码 tag `v0.18.0` 构建；可通过 `--vllm_ascend_version` 配置；会初始化 submodule；构建与运行依赖以华为云 Ascend PyPI 为主索引、官方 PyPI 为备用索引 |
| FLA | 从 `fla-org/flash-linear-attention` 源码 checkout，默认 `main` 分支；可通过 `--fla_version` 指定分支或 release tag |
| Megatron-LM | 源码 checkout，默认分支 `v0.15.3`；可通过 `--megatron_branch` 配置 |
| MindSpeed | 源码 checkout，默认分支 `core_r0.15.3`；可通过 `--mindspeed_branch` 配置 |
| mcore-bridge | PyPI 上的最新发布版；Ascend 构建入口没有独立的版本配置参数 |
| ms-swift | 来自 `modelscope/ms-swift` 的源码 checkout，默认分支 `main`；可通过 `--swift_branch` 配置 |
| DeepSpeed | 安装满足 `deepspeed>=0.19` 的最新发布包；`TORCH_DEVICE_BACKEND_AUTOLOAD=0` 仅作用于构建时的安装命令 |
| ModelScope | 通过 `pip install -U modelscope` 安装 PyPI 最新发布包；Ascend 镜像不再 clone ModelScope 或 modelscope-hub 源码仓库 |
| triton-ascend | 在 vLLM 之前安装，并在所有 editable 安装完成后再次固定版本；安装时以华为云 Ascend PyPI 为主索引、官方 PyPI 为备用索引；CANN `8.5.*` 默认 `3.2.0`、CANN `9.0.*` 默认 `3.2.1`；可通过 `--triton_ascend_version` 配置 |
```
## 镜像 Tag 说明
通过 `docker/build_image.py --image_type ascend` 构建的镜像使用以下 tag 格式：
已发布 tag 索引见 [`docker/ascend/supported_tags.md`](./supported_tags.md)。
```text
${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<TorchNPU-version>-<hardware-tag>-<os-tag>-<python-tag>-<arch>
```
```text
| 字段 | 示例 | 说明 |
|--------------------|----------------------------------|------|
| `swift-branch` | `main` | 构建镜像时使用的 ms-swift 分支 |
| `cann-version-tag` | `cann8.5.1`、`cann9.0.0` | 从 CANN 基础镜像 tag 解析 |
| `TorchNPU-version` | `2.9.0.post2` | 来自 `--torch_npu_version`，默认 `2.9.0.post2` |
| `hardware-tag` | `910b`、`a3`、`950` | 直接从 CANN 基础镜像 tag 解析；当前对应的 Atlas 平台名称为 `a2`、`a3`、`a5`，仅用于文档说明，不参与 tag 映射。 |
| `os-tag` | `ubuntu22.04`、`openeuler24.03` | 从 CANN 基础镜像 tag 解析；避免不同 OS 的镜像 tag 冲突 |
| `python-tag` | `py3.11` | 从 CANN 基础镜像 tag 解析 |
| `arch` | `aarch64`、`x86_64` | 从宿主机架构或 `--arch` 推导 |
```
A2 / CANN 9.0.0 示例（CANN 硬件 tag 为 `910b`）：
```text
${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-910b-ubuntu22.04-py3.11-aarch64
```
## 最新镜像
最新的 A2 和 A3 镜像发布在 [quay.io/ascend/ms-swift](https://quay.io/repository/ascend/ms-swift?tab=tags)。
**设备 / CANN 基础镜像 / OS / 镜像 Tag / Dockerfile**
- A3 — 9.1.0 — openEuler 24.03 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-openeuler24.03-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A3 — 9.1.0 — Ubuntu 22.04 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-ubuntu22.04-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A2 — 9.1.0 — openEuler 24.03 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-openeuler24.03-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
- A2 — 9.1.0 — Ubuntu 22.04 — `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-ubuntu22.04-py3.12` — [Dockerfile.ascend](./Dockerfile.ascend)
## 本地构建
先设置目标镜像仓库。构建脚本会把 `docker/ascend/Dockerfile.ascend` 渲染成根目录 `Dockerfile`，然后执行构建；Ascend 镜像分支当前不执行 push。
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend
```
完整固定版本参考（CANN 9.1.0、Atlas A2、Ubuntu 22.04、Python 3.12、ARM64）：
```bash
export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift
python docker/build_image.py \
--image_type ascend \
--base_image quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12 \
--soc_version ascend910b1 \
--arch arm \
--torch_version 2.10.0 \
--torch_npu_version 2.10.0.post2 \
--torchvision_version 0.25.0 \
--torchaudio_version 2.10.0 \
--vllm_version v0.23.0 \
--vllm_ascend_version v0.23.0 \
--fla_version main \
--triton_ascend_version 3.2.2 \
--modelscope_branch master \
--swift_branch v4.5.2 \
--megatron_branch core_v0.16.0 \
--mindspeed_branch core_r0.16.0
```
## 自定义构建参数
使用 `--image_type ascend` 选择 Ascend 构建器。当前 Ascend 构建路径实际使用的参数如下：
- `--base_image`（默认：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`）：选择 CANN、硬件、OS 和 Python；tag 必须符合 `<cann-version>-<hardware>-<os>-py<python-version>` 格式。
- `--soc_version`（默认：`ascend910_9391`）：设置 vLLM Ascend 构建的目标 SoC 和运行时 `SOC_VERSION`；必须与目标硬件匹配。[vLLM Ascend v0.23.0 安装文档](https://docs.vllm.ai/projects/ascend/en/v0.23.0/installation.html)给出了以下参考值。这些是 vLLM Ascend 上游的构建目标，不代表本 Dockerfile 已验证全部硬件路径：
- Atlas A2：`ascend910b1`
- Atlas A3：`ascend910_9391`
- Atlas 300I DUO：`ascend310p1`
- Atlas 950DT：`ascend950dt_9582`
- `--arch`（默认：根据宿主机自动检测）：可选 `arm` 或 `x86`；输出 tag 中会规范化为 `aarch64` 或 `x86_64`。
- `--torch_version`（默认：`2.9.0`）：选择 PyTorch；覆盖时必须同时传入匹配的 `--torchvision_version` 和 `--torchaudio_version`。
- `--torch_npu_version`（默认：`2.9.0.post2`）：基础版本必须与 `--torch_version` 完全一致。
- `--torchvision_version`（默认：`0.24.0`）：选择 torchvision；覆盖 `--torch_version` 时必须显式传入。
- `--torchaudio_version`（默认：`2.9.0`）：选择 torchaudio；覆盖 `--torch_version` 时必须显式传入。
- `--vllm_version`（默认：`0.18.0`）：选择官方 vLLM 源码 tag；可以带或不带前置 `v`。
- `--vllm_ascend_version`（默认：`0.18.0`）：选择官方 vLLM Ascend 源码 tag；可以带或不带前置 `v`。
- `--fla_version`（默认：`main`）：选择 FLA 分支或 release tag。
- `--triton_ascend_version`（默认值取决于 CANN 版本）：CANN 8.5 使用 `3.2.0`，CANN 9.0 使用 `3.2.1`；其他 CANN 系列（包括 9.1）必须显式传入。
- `--pip_extra_index_url`（默认：华为云 Ascend PyPI）：设置 Ascend 特定依赖的主索引。
- `--pypi_official_index_url`（默认：`https://pypi.org/simple`）：设置备用包索引。
- `--swift_branch`（默认：`main`）：选择 ms-swift 源码分支或 tag，并写入输出镜像 tag。
- `--megatron_branch`（默认：`v0.15.3`）：选择 Megatron-LM 源码分支或 tag。
- `--mindspeed_branch`（默认：`core_r0.15.3`）：选择 MindSpeed 源码分支或 tag。
Ascend 镜像的 Python 版本必须通过 `--base_image` 选择，`--python_version`
不会覆盖它。`--modelscope_branch` 虽然会被共用参数解析器接受，但 Ascend
Dockerfile 不使用该参数；当前通过 `pip install -U modelscope` 安装 ModelScope
最新发布包。
## 运行 Ascend 容器
宿主机需要提前安装兼容的 Ascend driver 和 firmware。容器通过挂载宿主机 NPU 设备和 driver 库使用昇腾硬件。
```bash
docker run --rm -it \
--name ms_swift_ascend \
--device /dev/davinci0 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \
-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-cann9.0.0-torch_npu2.9.0.post2-910b-ubuntu22.04-py3.11-aarch64 \
bash
```
进入容器后可以验证 NPU 和 Python 包：
```bash
npu-smi info
python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)"
python -c "import vllm, vllm_ascend; print('vllm ascend ok')"
pip show ms-swift modelscope mcore-bridge torch-npu triton-ascend
```
## 环境变量
```text
| 变量 | 值 |
|----------------------------|----|
| `SOC_VERSION` | 选定的 Ascend SoC，例如 `ascend910b1` 或 `ascend910_9391`；可通过 `--soc_version` 配置；用于指定 vLLM Ascend 的构建目标，并保留在运行时环境中（tag 中的硬件字段来自 `--base_image`） |
| `CANN_VERSION` | 从基础镜像 tag 解析得到 |
| `MEGATRON_LM_PATH` | `/Megatron-LM` |
| `PYTHONPATH` | 包含 `/Megatron-LM` |
| `VLLM_USE_MODELSCOPE` | `True` |
| `LMDEPLOY_USE_MODELSCOPE` | `True` |
| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` |
```
## 注意事项
- CANN、firmware 和 driver 版本必须互相兼容。
- Ubuntu 基础镜像通过 `apt-get` 安装系统依赖；openEuler 基础镜像通过 `yum` 安装对应 RPM 包。
- `triton-ascend` 和 vLLM Ascend editable 安装的依赖以 `https://mirrors.huaweicloud.com/ascend/repos/pypi` 为主索引、`https://pypi.org/simple` 为备用索引。pip 会汇总所有已配置索引的候选版本，并不保证严格的仓库优先级；请选择与 CANN、Python 和架构兼容的版本。
- 该镜像面向 Ascend NPU 上的 ms-swift 工作流。依赖安装过程中引入且与 NPU runtime 冲突的 CUDA-only 包会被移除。
- 生产任务建议使用固定镜像 tag，不要依赖浮动分支名。
## License
ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MindSpeed、TorchNPU、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。

41
docker/ascend/supported_tags.md Normal file
View File

@@ -0,0 +1,41 @@
# Supported Ascend Tags
This index lists the active tags published in `quay.io/ascend/ms-swift`. Tags
are grouped by the CANN version encoded in the tag. The `Dockerfile` column
points to the current Ascend build template; legacy tag formats are retained
for reference. The `content` column lists the software included in each image.
## swift 4.5.2
| Tag | Dockerfile | content |
| --- | --- | --- |
| `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-openeuler24.03-py3.12` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-a3-ubuntu22.04-py3.12` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-openeuler24.03-py3.12` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.5.2-cann9.1.0-torch_npu2.10.0.post2-910b-ubuntu22.04-py3.12` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
## swift 4.4.2
| Tag | Dockerfile | content |
| --- | --- | --- |
| `v4.4.2-cann9.0.1-torch_npu2.9.0.post4-a3-openeuler24.03-py3.11` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.4.2-cann9.0.1-torch_npu2.9.0.post4-a3-ubuntu22.04-py3.11` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.4.2-cann9.0.1-torch_npu2.9.0.post4-910b-openeuler24.03-py3.11` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.4.2-cann9.0.1-torch_npu2.9.0.post4-910b-ubuntu22.04-py3.11` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
## swift 4.3.0
| Tag | Dockerfile | content |
| --- | --- | --- |
| `v4.3.0-A3-py311-CANN9.0.0-ubuntu22.04` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.3.0-A2-py311-CANN9.0.0-ubuntu22.04` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
## swift 4.1.3
The following historical tags use the same software stack; their tag format
does not encode the full image metadata.
| Tag | Dockerfile | content |
| --- | --- | --- |
| `v4.1.3-A3-py311` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |
| `v4.1.3-A2-py311` | [Dockerfile.ascend](./Dockerfile.ascend) | toolkit/ops/nnal/torch_npu/ms-swift/vllm-ascend/mindspeed |

74
docker/build_image.py
View File

@@ -854,11 +854,16 @@ class AscendImageBuilder(StableGPUImageBuilder):
_DEFAULT_TORCH_NPU_VERSION = '2.9.0.post2' _DEFAULT_TORCH_NPU_VERSION = '2.9.0.post2'
_DEFAULT_VLLM_VERSION = '0.18.0' _DEFAULT_VLLM_VERSION = '0.18.0'
_DEFAULT_VLLM_ASCEND_VERSION = '0.18.0' _DEFAULT_VLLM_ASCEND_VERSION = '0.18.0'
_DEFAULT_FLA_VERSION = 'main'
_DEFAULT_PIP_EXTRA_INDEX_URL = (
'https://mirrors.huaweicloud.com/ascend/repos/pypi')
_DEFAULT_PYPI_OFFICIAL_INDEX_URL = 'https://pypi.org/simple'
_DEFAULT_TRITON_ASCEND_VERSIONS = { _DEFAULT_TRITON_ASCEND_VERSIONS = {
'8.5': '3.2.0', '8.5': '3.2.0',
'9.0': '3.2.1', '9.0': '3.2.1',
} }
_CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$') _CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$')
_HARDWARE_TAG_PATTERN = re.compile(r'^[0-9A-Za-z]+$')
_OS_TAG_PATTERN = re.compile(r'^[A-Za-z]+[0-9][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) _PYTHON_TAG_PATTERN = re.compile(r'^py\d+\.\d+$', re.IGNORECASE)
_TORCH_NPU_VERSION_PATTERN = re.compile( _TORCH_NPU_VERSION_PATTERN = re.compile(
@@ -881,25 +886,8 @@ class AscendImageBuilder(StableGPUImageBuilder):
'Please pass --arch x86 or --arch arm.') 'Please pass --arch x86 or --arch arm.')
return arch_mapping[arch] return arch_mapping[arch]
@staticmethod
def _get_atlas_hardware(soc_version: str) -> str:
soc_version = soc_version.lower()
atlas_mapping = {
'ascend910b1': 'A2',
'ascend910_9391': 'A3',
'ascend310p1': '300I',
}
if soc_version.startswith('ascend950'):
return 'A5'
if soc_version not in atlas_mapping:
raise ValueError(
f'Unsupported soc_version: {soc_version}. '
'Supported values are ascend910b1, ascend910_9391, '
'ascend310p1, and values starting with ascend950.')
return atlas_mapping[soc_version]
@classmethod @classmethod
def _get_cann_os_tags(cls, base_image: str) -> tuple: def _get_cann_image_tags(cls, base_image: str) -> tuple:
if ':' not in base_image.rsplit('/', 1)[-1]: if ':' not in base_image.rsplit('/', 1)[-1]:
raise ValueError( raise ValueError(
f'Ascend base image must include a tag: {base_image}') f'Ascend base image must include a tag: {base_image}')
@@ -913,11 +901,15 @@ class AscendImageBuilder(StableGPUImageBuilder):
f'{base_tag}') f'{base_tag}')
cann_version = parts[0] cann_version = parts[0]
hardware_tag = parts[1]
os_tag = parts[2] os_tag = parts[2]
python_tag = parts[3] python_tag = parts[3]
if not cls._CANN_VERSION_PATTERN.fullmatch(cann_version): if not cls._CANN_VERSION_PATTERN.fullmatch(cann_version):
raise ValueError(f'Invalid CANN version in Ascend base image tag: ' raise ValueError(f'Invalid CANN version in Ascend base image tag: '
f'{cann_version}') f'{cann_version}')
if not cls._HARDWARE_TAG_PATTERN.fullmatch(hardware_tag):
raise ValueError(f'Invalid hardware tag in Ascend base image tag: '
f'{hardware_tag}')
if not cls._OS_TAG_PATTERN.fullmatch(os_tag): if not cls._OS_TAG_PATTERN.fullmatch(os_tag):
raise ValueError( raise ValueError(
f'Invalid OS tag in Ascend base image tag: {os_tag}') f'Invalid OS tag in Ascend base image tag: {os_tag}')
@@ -925,7 +917,8 @@ class AscendImageBuilder(StableGPUImageBuilder):
raise ValueError( raise ValueError(
f'Invalid Python tag in Ascend base image tag: {python_tag}') f'Invalid Python tag in Ascend base image tag: {python_tag}')
return cann_version, f'CANN{cann_version}', os_tag, python_tag return (cann_version, f'CANN{cann_version}', hardware_tag, os_tag,
python_tag)
@staticmethod @staticmethod
def _get_os_family(os_tag: str) -> str: def _get_os_family(os_tag: str) -> str:
@@ -983,6 +976,7 @@ class AscendImageBuilder(StableGPUImageBuilder):
args.vllm_git_ref = cls._get_vllm_git_ref(args.vllm_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_git_ref = cls._get_vllm_git_ref(
args.vllm_ascend_version) args.vllm_ascend_version)
args.fla_version = args.fla_version or cls._DEFAULT_FLA_VERSION
if not args.triton_ascend_version: if not args.triton_ascend_version:
cann_series = '.'.join(args.cann_version.split('.')[:2]) cann_series = '.'.join(args.cann_version.split('.')[:2])
@@ -1002,12 +996,16 @@ class AscendImageBuilder(StableGPUImageBuilder):
if not args.base_image: if not args.base_image:
# Reuse the prebuilt vllm-ascend image to avoid rebuilding its stack. # 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' args.base_image = 'quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11'
args.pip_extra_index_url = (
args.pip_extra_index_url or self._DEFAULT_PIP_EXTRA_INDEX_URL)
args.pypi_official_index_url = (
args.pypi_official_index_url
or self._DEFAULT_PYPI_OFFICIAL_INDEX_URL)
self._init_torch_versions(args) self._init_torch_versions(args)
args.arch = self._normalize_arch(args.arch) 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.hardware_tag,
(args.cann_version, args.cann_version_tag, args.os_tag, args.os_tag, args.ascend_python_tag) = (
args.ascend_python_tag) = ( self._get_cann_image_tags(args.base_image))
self._get_cann_os_tags(args.base_image))
self._get_os_family(args.os_tag) self._get_os_family(args.os_tag)
self._init_component_versions(args) self._init_component_versions(args)
return super().init_args(args) return super().init_args(args)
@@ -1016,13 +1014,13 @@ class AscendImageBuilder(StableGPUImageBuilder):
return self.args.ascend_python_tag return self.args.ascend_python_tag
def generate_dockerfile(self) -> str: def generate_dockerfile(self) -> str:
extra_content = """ with open('docker/ascend/Dockerfile.ascend', 'r') as f:
RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
"""
with open('docker/Dockerfile.ascend', 'r') as f:
content = f.read() content = f.read()
content = content.replace('{base_image}', self.args.base_image) content = content.replace('{base_image}', self.args.base_image)
content = content.replace('{pip_extra_index_url}',
self.args.pip_extra_index_url)
content = content.replace('{pypi_official_index_url}',
self.args.pypi_official_index_url)
content = content.replace('{soc_version}', self.args.soc_version) content = content.replace('{soc_version}', self.args.soc_version)
content = content.replace('{cann_version}', self.args.cann_version) content = content.replace('{cann_version}', self.args.cann_version)
content = content.replace('{torch_version}', content = content.replace('{torch_version}',
@@ -1038,11 +1036,8 @@ RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
self.args.vllm_ascend_git_ref) self.args.vllm_ascend_git_ref)
content = content.replace('{triton_ascend_version}', content = content.replace('{triton_ascend_version}',
self.args.triton_ascend_version) self.args.triton_ascend_version)
content = content.replace('{extra_content}', extra_content) content = content.replace('{fla_version}', self.args.fla_version)
content = content.replace('{cur_time}', formatted_time) content = content.replace('{cur_time}', formatted_time)
content = content.replace('{install_ms_deps}', 'False')
content = content.replace('{modelscope_branch}',
self.args.modelscope_branch)
content = content.replace('{swift_branch}', self.args.swift_branch) content = content.replace('{swift_branch}', self.args.swift_branch)
content = content.replace('{megatron_branch}', content = content.replace('{megatron_branch}',
self.args.megatron_branch) self.args.megatron_branch)
@@ -1053,7 +1048,7 @@ RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
def image(self) -> str: def image(self) -> str:
tag = (f'{self.args.swift_branch}-{self.args.cann_version_tag}-' tag = (f'{self.args.swift_branch}-{self.args.cann_version_tag}-'
f'torch_npu{self.args.torch_npu_version}-' f'torch_npu{self.args.torch_npu_version}-'
f'{self.args.atlas_hardware}-{self.args.os_tag}-' f'{self.args.hardware_tag}-{self.args.os_tag}-'
f'{self.args.python_tag}-' f'{self.args.python_tag}-'
f'{self.args.arch}') f'{self.args.arch}')
return f'{docker_registry}:{tag.lower()}' return f'{docker_registry}:{tag.lower()}'
@@ -1076,7 +1071,20 @@ parser.add_argument('--torchaudio_version', type=str, default=None)
parser.add_argument('--tf_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_version', type=str, default=None)
parser.add_argument('--vllm_ascend_version', type=str, default=None) parser.add_argument('--vllm_ascend_version', type=str, default=None)
parser.add_argument('--fla_version', type=str, default=None)
parser.add_argument('--triton_ascend_version', type=str, default=None) parser.add_argument('--triton_ascend_version', type=str, default=None)
parser.add_argument(
'--pip_extra_index_url',
type=str,
default=None,
help='Primary Python package index for Ascend-specific installs. Defaults '
'to the Huawei Cloud Ascend PyPI repository for Ascend images.')
parser.add_argument(
'--pypi_official_index_url',
type=str,
default=None,
help='Fallback Python package index for Ascend-specific installs. Defaults '
'to the official PyPI repository.')
parser.add_argument('--lmdeploy_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('--flashattn_version', type=str, default=None)
parser.add_argument('--autogptq_version', type=str, default=None) parser.add_argument('--autogptq_version', type=str, default=None)