GeezerVison

# ---------- 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 clone --depth 1 --branch {megatron_adaptor_branch} https://gitcode.com/Ascend/MegatronAdaptor.git /MegatronAdaptor && \

git clone --depth 1 --branch {transformer_engine_npu_branch} https://gitcode.com/Ascend/TransformerEngineNPU.git /TransformerEngineNPU && \

git clone --depth 1 --single-branch --branch {mcore_bridge_branch} 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

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

cd /mcore-bridge && pip install --no-cache-dir -e . && \

cd /MegatronAdaptor && pip install --no-cache-dir -e . && \

cd /TransformerEngineNPU && pip install --no-cache-dir -e . --no-build-isolation && \

pip install --no-cache-dir -U modelscope && \

cd /ms-swift && pip install --no-cache-dir -e .

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 -U qwen_vl_utils qwen_omni_utils 'librosa<1.0' 'timm>=0.9.0' transformers accelerate 'peft>=0.11,<0.21' 'trl>=0.15,<1.0' 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 && \

# ---------- Re-pin triton-ascend after all editable installs ----------

# vllm-ascend may otherwise replace the CANN-specific version selected above.

# Restore NumPy 2 for PyArrow after Triton-Ascend pins NumPy 1.26.4.

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}

triton-ascend=={triton_ascend_version} && \

pip install --no-cache-dir --upgrade --no-deps 'numpy>=2.0'

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

# 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)"

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', 'megatron_adaptor': '/MegatronAdaptor', 'transformer_engine': '/TransformerEngineNPU', 'mcore_bridge': '/mcore-bridge', '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('mindspeed') is None; 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"]

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.

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, MegatronAdaptor, TransformerEngineNPU, mcore-bridge, ms-swift, and ModelScope runtime components needed for Ascend inference and training workflows.

--image_type ascend

Complete version-pinned reference (CANN 9.1.0, Atlas A2, Ubuntu 22.04,

Python 3.12, ARM64):

Complete build reference (CANN 9.1.0, Atlas A2, Ubuntu 22.04,

Python 3.12, ARM64). The Megatron components follow the [ms-swift NPU installation guide](https://swift.readthedocs.io/zh-cn/latest/BestPractices/NPU-support.html#megatronadaptor-transformerenginenpu): MCore 0.18, MegatronAdaptor and TransformerEngineNPU, without MindSpeed. Use compatible ms-swift and mcore-bridge branches containing this integration; the example uses `main`.

export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift

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

--mcore_bridge_branch v1.6.2

--swift_branch main \

--megatron_branch core_r0.18.0 \

--megatron_adaptor_branch core_r0.18.0 \

--transformer_engine_npu_branch main \

--mcore_bridge_branch main

- `--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.

- `--megatron_branch` (default: `core_r0.18.0`): selects the Megatron-LM source branch or tag.

- `--megatron_adaptor_branch` (default: `core_r0.18.0`): selects the MegatronAdaptor source branch or tag.

- `--transformer_engine_npu_branch` (default: `main`): selects the TransformerEngineNPU source branch or tag.

- `--mcore_bridge_branch` (default: `main`): selects the mcore-bridge source branch or release tag for the editable install.

`--python_version` does not override the Python version for Ascend images;

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.

ms-swift and ModelScope components follow their upstream repository licenses. CANN, MegatronAdaptor, TransformerEngineNPU, TorchNPU, vLLM Ascend, and other pre-installed third-party components are subject to their own upstream licenses.

ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms-swift 运行环境。镜像基于 Ascend CANN 容器镜像构建，包含 Ascend 推理和训练工作流所需的 Python、CANN、TorchNPU、vLLM Ascend、FLA、Megatron、MindSpeed、mcore-bridge、ms-swift 以及 ModelScope 运行组件。

ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms-swift 运行环境。镜像基于 Ascend CANN 容器镜像构建，包含 Ascend 推理和训练工作流所需的 Python、CANN、TorchNPU、vLLM Ascend、FLA、Megatron、MegatronAdaptor、TransformerEngineNPU、mcore-bridge、ms-swift 以及 ModelScope 运行组件。

--image_type ascend

完整固定版本参考（CANN 9.1.0、Atlas A2、Ubuntu 22.04、Python 3.12、ARM64）：

完整构建参考（CANN 9.1.0、Atlas A2、Ubuntu 22.04、Python 3.12、ARM64）。Megatron 组件遵循 [ms-swift NPU 安装文档](https://swift.readthedocs.io/zh-cn/latest/BestPractices/NPU-support.html#megatronadaptor-transformerenginenpu)：使用 MCore 0.18、MegatronAdaptor 和 TransformerEngineNPU，不安装 MindSpeed。ms-swift 和 mcore-bridge 应选择已包含配套适配的兼容分支；示例使用 `main`。

export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift

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

--mcore_bridge_branch v1.6.2

--swift_branch main \

--megatron_branch core_r0.18.0 \

--megatron_adaptor_branch core_r0.18.0 \

--transformer_engine_npu_branch main \

--mcore_bridge_branch main

- `--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。

- `--megatron_branch`（默认：`core_r0.18.0`）：选择 Megatron-LM 源码分支或 tag。

- `--megatron_adaptor_branch`（默认：`core_r0.18.0`）：选择 MegatronAdaptor 源码分支或 tag。

- `--transformer_engine_npu_branch`（默认：`main`）：选择 TransformerEngineNPU 源码分支或 tag。

- `--mcore_bridge_branch`（默认：`main`）：选择 mcore-bridge 源码分支或 release tag，并以 editable 模式安装。

Ascend 镜像的 Python 版本必须通过 `--base_image` 选择，`--python_version`

ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MindSpeed、TorchNPU、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。

ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MegatronAdaptor、TransformerEngineNPU、TorchNPU、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。

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)

content = content.replace('{megatron_adaptor_branch}',

self.args.megatron_adaptor_branch)

content = content.replace('{transformer_engine_npu_branch}',

self.args.transformer_engine_npu_branch)

content = content.replace('{mcore_bridge_branch}',

self.args.mcore_bridge_branch)

return content

help='Minor build only: install modelscope from the released PyPI version '

'(--modelscope_version, default) or from --modelscope_branch source.')

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('--megatron_branch', type=str, default='core_r0.18.0')

parser.add_argument(

'--megatron_adaptor_branch', type=str, default='core_r0.18.0')

parser.add_argument(

'--transformer_engine_npu_branch', type=str, default='main')

parser.add_argument('--mcore_bridge_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)

# An empty string against "ab" used to return distance 2, so the

# score was -1. The normalized distance is 1, so the score is 0.

# cat/cut is still 2/3 and an exact match is still 1.

self.assertAlmostEqual(score, (0.0 + 0.0 + (1.0 - 1.0 / 3.0) + 1.0) / 4.0)

expected_score = (0.0 + 0.0 + (1.0 - 1.0 / 3.0) + 1.0) / 4.0

self.assertAlmostEqual(score, expected_score)

if __name__ == '__main__':