Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-07-09 20:09:17 +02:00
Code Issues Packages Projects Releases Wiki Activity

Update npu dockerfile (#1736)

Browse Source
This commit is contained in:
addsubmuldiv
2026-06-15 11:07:56 +08:00
committed by GitHub
parent c7dba8de1e
commit c1d88e6ab7
4 changed files with 389 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
docker/Dockerfile.ascend
View File

@@ -3,7 +3,8 @@ FROM {base_image}
ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_DEFAULT_TIMEOUT=300 \
PIP_RETRIES=10 \
SOC_VERSION={soc_version}
SOC_VERSION={soc_version} \
CANN_VERSION={cann_version}
SHELL ["/bin/bash", "-c"]
@@ -36,15 +37,15 @@ 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 torchvision==0.24.0 && \
pip install torch==2.9.0 torch_npu==2.9.0.post2 torchvision==0.24.0 && \
# 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 v0.15.3 https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \
git clone --depth 1 --branch core_r0.15.3 https://gitcode.com/Ascend/MindSpeed.git /MindSpeed && \
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
@@ -70,7 +71,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \
fi && \
pip install --no-cache-dir --force-reinstall --no-deps \
torch_npu==2.9.0 && \
torch_npu==2.9.0.post2 && \
rm -rf /root/.cache/pip
# ---------- Remove CUDA-only dependencies pulled in by vllm (they cause missing libtorch_cuda.so errors on NPU) ----------
@@ -122,6 +123,37 @@ 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 ----------
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
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

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

@@ -0,0 +1,156 @@
# 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`
- Default output tag: `${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-<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` |
| 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` |
| 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` |
| 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` |
## 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>
```
| 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` |
Default example on an ARM64 host:
```text
${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm
```
A2 / CANN 9.0.0 example:
```text
${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm
```
## 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
```
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
```
For slow networks, Linux hosts can use Docker host networking after the root `Dockerfile` is generated:
```bash
docker build --network host \
-t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
-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-A2-py311-CANN9.0.0-ubuntu22.04-arm \
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.
- CANN `8.5.*` and CANN `9.0.0` use different `triton-ascend` install paths in this Dockerfile.
- 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.

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

@@ -0,0 +1,156 @@
# 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`
- 默认输出 tag：`${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-<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-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` |
| Megatron-LM | 源码 checkout，默认分支 `v0.15.3` |
| MindSpeed | 源码 checkout，默认分支 `core_r0.15.3` |
| mcore-bridge | 来自 `modelscope/mcore-bridge` 的源码 checkout |
| 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 |
## 支持的 Tag 格式
通过 `docker/build_image.py --image_type ascend` 构建的镜像使用以下 tag 格式：
```text
${DOCKER_REGISTRY}:<swift-branch>-<atlas-hardware>-<python-tag>-<cann-version-tag>-<os-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` 推导 |
ARM64 宿主机上的默认示例：
```text
${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm
```
A2 / CANN 9.0.0 示例：
```text
${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm
```
## 本地构建
先设置目标镜像仓库。构建脚本会把 `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
```
需要时可以覆盖 Megatron 或 MindSpeed 源码分支：
```bash
python docker/build_image.py \
--image_type ascend \
--megatron_branch v0.15.3 \
--mindspeed_branch core_r0.15.3
```
如果构建时网络较慢，Linux 宿主机可以在根目录 `Dockerfile` 生成后使用 host network 构建：
```bash
docker build --network host \
-t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
-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-A2-py311-CANN9.0.0-ubuntu22.04-arm \
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 版本必须互相兼容。
- 这个 Dockerfile 对 CANN `8.5.*` 和 CANN `9.0.0` 使用不同的 `triton-ascend` 安装路径。
- 该镜像面向 Ascend NPU 上的 ms-swift 工作流。依赖安装过程中引入且与 NPU runtime 冲突的 CUDA-only 包会被移除。
- 生产任务建议使用固定镜像 tag，不要依赖浮动分支名。
## License
ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MindSpeed、torch-npu、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。

41
docker/build_image.py
View File

@@ -1,6 +1,7 @@
import argparse
import os
import platform
import re
import subprocess
from datetime import datetime
from typing import Any
@@ -462,6 +463,9 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
class AscendImageBuilder(StableGPUImageBuilder):
_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.]*$')
@staticmethod
def _normalize_arch(arch: str = None) -> str:
arch = arch or platform.machine()
@@ -496,12 +500,39 @@ class AscendImageBuilder(StableGPUImageBuilder):
'ascend310p1, and values starting with ascend950.')
return atlas_mapping[soc_version]
@classmethod
def _get_cann_os_tags(cls, base_image: str) -> tuple:
if ':' not in base_image.rsplit('/', 1)[-1]:
raise ValueError(
f'Ascend base image must include a tag: {base_image}')
base_tag = base_image.rsplit(':', 1)[1]
parts = base_tag.split('-')
if len(parts) < 4:
raise ValueError(
'Ascend base image tag must look like '
f'<cann_version>-<hardware>-<os_tag>-py<version>, got: '
f'{base_tag}')
cann_version = parts[0]
os_tag = parts[2]
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}')
return cann_version, f'CANN{cann_version}', os_tag
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'
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))
return super().init_args(args)
def generate_dockerfile(self) -> str:
@@ -512,18 +543,24 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
content = f.read()
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('{extra_content}', extra_content)
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('{megatron_branch}',
self.args.megatron_branch)
content = content.replace('{mindspeed_branch}',
self.args.mindspeed_branch)
return content
def image(self) -> str:
return (
f'{docker_registry}:{self.args.swift_branch}-'
f'{self.args.atlas_hardware}-{self.args.python_tag}-{self.args.arch}'
f'{self.args.atlas_hardware}-{self.args.python_tag}-'
f'{self.args.cann_version_tag}-{self.args.os_tag}-{self.args.arch}'
)
def push(self):
@@ -549,6 +586,8 @@ parser.add_argument('--optimum_version', type=str, default=None)
parser.add_argument('--modelscope_branch', type=str, default='master')
parser.add_argument('--modelscope_version', type=str, default='9.99.0')
parser.add_argument('--swift_branch', type=str, default='main')
parser.add_argument('--megatron_branch', type=str, default='v0.15.3')
parser.add_argument('--mindspeed_branch', type=str, default='core_r0.15.3')
parser.add_argument('--soc_version', type=str, default='ascend910_9391')
parser.add_argument('--arch', type=str, choices=['x86', 'arm'], default=None)
parser.add_argument('--dry_run', type=int, default=0)