Ascend NPU Quick Start#
Branch notice: Ascend NPU support is currently maintained on the ascend branch (not yet on
main), with plans to merge intomainlater. Clone or checkout that branch before running any NPU examples below.
⚠️ If you encounter problems running vime on Ascend NPU, feel free to open an issue on vllm-project/vime.
Overview#
vime on Ascend NPU uses the Megatron training backend together with the
vLLM Ascend rollout backend. In decoupled mode, actor weights sync to vLLM
over HCCL; in colocate mode (--colocate), weights sync over NPU IPC.
Current support targets Ascend Atlas A2 / A3 (aarch64) hardware.
Get the Ascend Branch#
git clone --branch ascend https://github.com/vllm-project/vime.git
cd vime
If you already have the repo:
git fetch origin ascend
git checkout ascend
Ascend Branch Resources#
Resource |
Description |
|---|---|
Full NPU guide with end-to-end GRPO example and training flags |
|
Source-build guide, pinned commits, and patch list |
|
Qwen3-4B decoupled training (4 actor + 4 rollout NPUs) |
|
Qwen3-30B-A3B MoE NPU training script |
|
Model args for Qwen3-30B-A3B on NPU |
Basic Environment Setup#
Docker Image#
The recommended path for validation is the published vime NPU image:
export IMAGE=quay.io/ascend/vime:vime-latest
# A2: export IMAGE=quay.io/ascend/vime:vime-a2-latest
docker pull "${IMAGE}"
For source builds and dependency debugging, follow
docker/npu_patch/README.md
on the ascend branch.
Pull and Start Docker Container#
Start the container with Ascend devices and driver files mounted. Device names and mount paths vary by host; reuse the mounts from a known working vLLM Ascend container if the layout differs.
docker run -d --name vime-npu -it --net=host --shm-size=1024g \
--privileged=true \
--cap-add=SYS_PTRACE \
--device=/dev/davinci_manager \
--device=/dev/hisi_hdc \
--device=/dev/devmm_svm \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin:/usr/local/sbin \
-v /home:/home \
-v /mnt:/mnt \
-v /tmp:/tmp \
-v /data:/data \
-v /path/to:/path/to \
-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
"${IMAGE}"
docker exec -it vime-npu bash
Inside the container, initialize the CANN environment before training:
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh
Model and Dataset Download#
export MODEL_ROOT=/root
mkdir -p ${MODEL_ROOT}/models ${MODEL_ROOT}/datasets
# Model weights (Qwen3-4B)
hf download Qwen/Qwen3-4B --local-dir ${MODEL_ROOT}/models/Qwen3-4B
# Training dataset (dapo-math-17k)
hf download --repo-type dataset zhuzilin/dapo-math-17k \
--local-dir ${MODEL_ROOT}/datasets/dapo-math-17k
Training (Qwen3-4B Example)#
After checking out the ascend branch inside the container, run the bundled
script:
cd /root/vime
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh
MODEL_ROOT=/root bash scripts/run-qwen3-4B-npu.sh
The full log is written to /root/vime/train_qwen3_4b_vllm.log.
Note: The main difference from the NVIDIA workflow is Ascend-specific environment variables — use
ASCEND_RT_VISIBLE_DEVICESinstead ofCUDA_VISIBLE_DEVICES, and setRAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1so Ray schedules NPUs correctly. The reference script targets an Atlas A3 host with 16 visible NPUs; on an 8-NPU host, setASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7.
For the full training command, HCCL port ranges, and flag explanations, see NPU.md on the ascend branch.
MoE Example (Qwen3-30B-A3B)#
For the MoE model on NPU, use the scripts on the ascend branch:
bash scripts/run-qwen3-30B-A3B-npu.sh
See scripts/models/qwen3-30B-A3B-npu.sh for model-specific arguments.