diff --git a/docker/Dockerfile.ascend b/docker/Dockerfile.ascend index b18eff2f..c9f4e80e 100644 --- a/docker/Dockerfile.ascend +++ b/docker/Dockerfile.ascend @@ -1,14 +1,13 @@ FROM {base_image} -# Use bash so that `source` and other bash builtins work in all following RUN steps. +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_DEFAULT_TIMEOUT=300 \ + PIP_RETRIES=10 \ + SOC_VERSION={soc_version} + SHELL ["/bin/bash", "-c"] -ENV PIP_DISABLE_PIP_VERSION_CHECK=1 -ENV PIP_DEFAULT_TIMEOUT=300 -ENV PIP_RETRIES=10 -ENV TRANSFORMERS_VERBOSITY=error -ENV TRANSFORMERS_NO_ADVISORY_WARNINGS=1 - +# ---------- 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 && \ @@ -19,23 +18,62 @@ RUN rm -f /etc/apt/apt.conf.d/docker-clean && \ rm -rf /var/lib/apt/lists/* RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ - pip config set install.trusted-host mirrors.aliyun.com + pip config set install.trusted-host mirrors.aliyun.com && \ + ARCH=$(uname -m) && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi {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 v0.14.0 https://github.com/vllm-project/vllm && \ + git clone --depth 1 --branch v0.14.0rc1 https://github.com/vllm-project/vllm-ascend.git -# Reuse the vllm-ascend base image and only add the extra repos we need. -# --depth 1 keeps the image smaller; branch/tag names work with shallow clone. +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 && \ + # 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 && \ 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 +# ---------- 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 cache purge + 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 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==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \ + else \ + pip install --no-cache-dir --force-reinstall --no-deps \ + 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 && \ + 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 @@ -68,6 +106,8 @@ fi ARG CUR_TIME={cur_time} RUN echo $CUR_TIME +RUN pip install --no-cache-dir --no-build-isolation OpenCC + 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 && \ pip install --no-cache-dir -U funasr scikit-learn && \ diff --git a/docker/build_image.py b/docker/build_image.py index 37d35787..a80879b0 100644 --- a/docker/build_image.py +++ b/docker/build_image.py @@ -1,5 +1,6 @@ import argparse import os +import platform import subprocess from datetime import datetime from typing import Any @@ -461,10 +462,46 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy class AscendImageBuilder(StableGPUImageBuilder): + @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', + } + if arch not in arch_mapping: + raise ValueError(f'Unsupported architecture: {arch}. ' + 'Please pass --arch x86 or --arch arm.') + 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] + 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/vllm-ascend:v0.14.0rc1-a3' + 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) return super().init_args(args) def generate_dockerfile(self) -> str: @@ -474,6 +511,7 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy with open('docker/Dockerfile.ascend', 'r') as f: content = f.read() content = content.replace('{base_image}', self.args.base_image) + content = content.replace('{soc_version}', self.args.soc_version) content = content.replace('{extra_content}', extra_content) content = content.replace('{cur_time}', formatted_time) content = content.replace('{install_ms_deps}', 'False') @@ -484,8 +522,9 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy def image(self) -> str: return ( - f'{docker_registry}:{self.args.base_image.split(":")[-1]}-torch2.7.1' - f'-{self.args.modelscope_version}-ascend-test') + f'{docker_registry}:{self.args.swift_branch}-' + f'{self.args.atlas_hardware}-{self.args.python_tag}-{self.args.arch}' + ) def push(self): return 0 @@ -510,6 +549,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('--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) args = parser.parse_args()