[RFC] Steps and Test Result for Running Qwen3-4B on NPU(A3) （multi-agent rollout pipeline）

[Windfeng8]

This issue documents the end-to-end deployment and training workflow for running Qwen3-4B GRPO training (multi-agent rollout pipeline)with VIME on Ascend NPU (Atlas 800I A3, 16× NPUs), intended as a reference for the community.

The reference guide PR is as follows:#298

Multi-Agent RL

This directory provides an example of running multi-agent reinforcement learning (RL) with vime.

New Arguments

• Specify the agent rollout function with the --custom-generate-function-path argument.
• Set the --rollout-max-context-len argument according to your model’s context window.

ROLLOUT_ARGS=(
   --custom-generate-function-path examples.multi_agent.rollout_with_multi_agents.generate_with_multi_agents
   --prompt-data "${DATA_FILE}"
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --rm-type math

   --rollout-backend vllm
   --vllm-weight-sync-mode native
   --vllm-gpu-memory-utilization 0.6
   --vllm-enable-sleep-mode
   --vllm-max-model-len 4096
   --vllm-enforce-eager

   --num-rollout 200
   --rollout-batch-size 32
   --n-samples-per-prompt 8
   --rollout-max-context-len 4096
   --rollout-max-response-len 2048
   --rollout-temperature 1.0

   --global-batch-size 256
   --balance-data
)

Preparing the Running Environment

Use docker

# Update the vime image
export IMAGE=quay.io/ascend/vime:vime-latest

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 /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
    $IMAGE

docker exec -it vime-npu bash

Use Source code

Python Version

Only python==3.12 is supported currently.

conda create -n vime_release python=3.12
conda activate vime_release

Working Directory Setup

mkdir <WORKSPACE> && cd <WORKSPACE>

CANN Environment

Prior to start work with vime on Ascend you need to install CANN Toolkit, Kernels operator package and NNAL version 9.0.0 , check the installation guide

source <CANN_PATH>/ascend-toolkit/set_env.sh
source <CANN_PATH>/nnal/atb/set_env.sh

PyTorch and PyTorch NPU

pip install torch-npu==2.10.0

Installing Dependencies

Megatron-Bridge

pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps

cd <WORKSPACE>
git clone https://github.com/radixark/Megatron-Bridge.git -b bridge
git checkout 3fd3768045422d0aa5c97e90a4e6c659aea9acb9
pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation

Megatron-LM

cd <WORKSPACE>
git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \
  cd Megatron-LM/ && git checkout 3714d81d418c9f1bca4594fc35f9e8289f652862 && \
  pip install -e .

MindSpeed

cd <WORKSPACE>
git clone https://gitcode.com/Ascend/MindSpeed.git && \
  cd MindSpeed/ && git checkout fc63de5c48426dd019c3b3f39e65f5bdf56e4086 && \
  pip install -e .

Vime

cd <WORKSPACE>
git clone https://github.com/vllm-project/vime.git -b fix/multi-agent && cd vime
cp -r docker/npu_patch ../npu_patch
pip install -e .

Applying Patches

cd <WORKSPACE>/Megatron-LM
git apply ../npu_patch/megatron_common.patch
git apply ../npu_patch/megatron.patch

cd <WORKSPACE>/Megatron-Bridge
git apply ../npu_patch/megatron-bridge.patch

cd <WORKSPACE>/MindSpeed
git apply ../npu_patch/mindspeed.patch

Additional Dependencies

cd <WORKSPACE>/vime
pip install triton-ascend
pip install torch-npu==2.10.0
pip install torchvision==0.25.0
pip install numpy==1.26.4

Training script

examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh

#!/bin/bash

# for rerun the task
pkill -9 -f "vllm serve"
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python

set -ex

# will prevent ray from buffering stdout/stderr
export PYTHONUNBUFFERED=1

export SLIME_SCRIPT_TRAIN_BACKEND=megatron
export PYTHONPATH="/workspace/issue205/Megatron-Bridge/src:/workspace/issue205/Megatron-Bridge:/workspace/issue205/Megatron-LM:/workspace/issue205/vime-ascend:${PYTHONPATH}"
export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}"
export CUDA_DEVICE_MAX_CONNECTIONS=1
export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
export HYDRA_FULL_ERROR=1
export MASTER_PORT="${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}"
export DISABLE_L2_CACHE=1
export VLLM_ASCEND_ENABLE_NZ=0

VIME_ROOT="${VIME_ROOT:-/workspace/issue205/vime-ascend}"
SCRIPT_DIR="${VIME_ROOT}/scripts"
WEIGHT_DIR="${WEIGHT_DIR:-/home/data/weights/Qwen3-4B}"
DATA_FILE="${DATA_FILE:-/home/w00893744/dataset/dapo-math-17k.jsonl}"
RUN_TS="${RUN_TS:-$(date +%Y%m%d_%H%M%S)}"
LOG_FILE="${LOG_FILE:-/home/w00893744/train_qwen3_4b_multi_agent_vllm_${RUN_TS}.log}"
MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"

cd "${VIME_ROOT}"
source "${SCRIPT_DIR}/models/qwen3-4B.sh"

CKPT_ARGS=(
   --hf-checkpoint "${WEIGHT_DIR}"
   --load "${WEIGHT_DIR}"
   --megatron-to-hf-mode bridge
)

ROLLOUT_ARGS=(
   --custom-generate-function-path examples.multi_agent.rollout_with_multi_agents.generate_with_multi_agents
   --prompt-data "${DATA_FILE}"
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --rm-type math

   --rollout-backend vllm
   --vllm-weight-sync-mode native
   --vllm-gpu-memory-utilization 0.6
   --vllm-enable-sleep-mode
   --vllm-max-model-len 4096
   --vllm-enforce-eager

   --num-rollout 200
   --rollout-batch-size 32
   --n-samples-per-prompt 8
   --rollout-max-context-len 4096
   --rollout-max-response-len 2048
   --rollout-temperature 1.0

   --global-batch-size 256
   --balance-data
)

EVAL_ARGS=(
)

PERF_ARGS=(
   --tensor-model-parallel-size 4
   --pipeline-model-parallel-size 1
   --context-parallel-size 1
   --expert-model-parallel-size 1
   --expert-tensor-parallel-size 1

   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --use-dynamic-batch-size
   --max-tokens-per-gpu 8192
   --micro-batch-size 1
)

GRPO_ARGS=(
   --advantage-estimator grpo
   --kl-loss-coef 0.0
   --kl-loss-type low_var_kl
   --kl-coef 0.00
   --entropy-coef 0.0
   --eps-clip 0.2
   --eps-clip-high 0.28
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-6
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)

WANDB_ARGS=(
)

VLLM_ARGS=(
   --rollout-num-gpus-per-engine 4
)

MISC_ARGS=(
   --attention-dropout 0.0
   --hidden-dropout 0.0
   --accumulate-allreduce-grads-in-fp32
   --attention-softmax-in-fp32
   --attention-backend flash
   --use-flash-attn
   --train-memory-margin-bytes 2147483648
)

# launch the master node of ray in container
export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265

python "${VIME_ROOT}/train.py" \
   --train-backend megatron \
   --actor-num-nodes 1 \
   --actor-num-gpus-per-node 4 \
   --rollout-num-gpus 4 \
   ${MODEL_ARGS[@]} \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${GRPO_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${EVAL_ARGS[@]} \
   ${VLLM_ARGS[@]} \
   ${MISC_ARGS[@]} \
   2>&1 | tee -a "${LOG_FILE}"

Result Comparison

The above shows the results of running 60 steps. Although the raw_reward increased, the corresponding time required also increased.