昇腾 NPU 端到端配置样例
最后更新:2026/04/27。
本文档提供在华为昇腾 NPU 上运行 ROLL 的端到端配置样例,涵盖环境准备、资源切分和启动命令,适用于单机和多机场景。
前置条件
运行本样例前,请确保:
- 已拉取或构建与硬件匹配的昇腾镜像(参见 Docker 使用指南)。
- 已在容器内验证环境(参见 验证环境)。
- 已将模型权重下载到容器可访问的目录。
当前仓库在 examples/ascend_examples 中提供可直接运行的昇腾 RLVR 示例,包括 qwen3_30b_rlvr_fsdp2.yaml 和 run_rlvr_pipeline.sh。
GPU 与 NPU 的关键差异
将 GPU 配置适配到 NPU 时,必须进行以下修改:
| 项目 | GPU | NPU |
|---|---|---|
| 训练后端 | Megatron 或 FSDP2 | 仅 FSDP2(NPU 不支持 Megatron) |
| 注意力实现 | flash_attn 或 fa2 | 通过 transformers 使用 fa2(不能使用 flash_attn 包) |
| 通信后端 | NCCL | HCCL |
| 设备可见性 | CUDA_VISIBLE_DEVICES | ASCEND_RT_VISIBLE_DEVICES |
样例 1:单机 Agentic 流水线(Qwen2.5-0.5B)
本样例在单个 8 卡 NPU 节点上使用 FSDP2 运行 FrozenLake Agentic 流水线。
步骤 1:启动容器
docker run -dit \
--name roll_npu_single \
--ulimit nofile=65536:65536 \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci6 \
--device /dev/davinci7 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/Ascend/add-ons:/usr/local/Ascend/add-ons \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /path/to/models:/data/models \
-v /path/to/data:/data \
--ipc=host \
--net=host \
roll:ascend-a3 \
/bin/bash
步骤 2:设置环境变量
# HCCL 通信
export HCCL_CONNECT_TIMEOUT=3600
export HCCL_DETERMINISTIC=false
export HCCL_OP_EXPANSION_MODE="AIV"
export HCCL_NPU_SOCKET_PORT_RANGE="auto"
# NPU 显存
export NPU_MEMORY_FRACTION=0.96
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export MULTI_STREAM_MEMORY_REUSE=1
export TASK_QUEUE_ENABLE=2
export COMBINED_ENABLE=1
# CPU 调度
export CPU_AFFINITY_CONF=2
export OMP_NUM_THREADS=1
# vLLM-Ascend 推理
export VLLM_USE_V1=1
export VLLM_ASCEND_ENABLE_NZ=0
export VLLM_ASCEND_ENABLE_FLASHCOMM=1
export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
# 算子编译缓存
export ACL_OP_COMPILER_CACHE_MODE=enable
export ACL_OP_COMPILER_CACHE_DIR=/tmp/npu_cache
export ASCEND_MAX_OP_CACHE_SIZE=5000
# 日志(生产环境)
export ASCEND_GLOBAL_LOG_LEVEL=3
export ASDOPS_LOG_LEVEL=ERROR
export ATB_LOG_LEVEL=ERROR
步骤 3:创建 NPU 配置文件
创建 YAML 配置文件 examples/agentic_frozen_lake_npu/agentic_frozen_lake_npu.yaml,以下为 NPU 专用配置。该路径很重要:配置中的 defaults: - ../config/traj_envs@_here_ 依赖 examples/ 下的相对目录结构;如果放到其他目录,需要同步调整 defaults 路径。与 GPU 配置的关键差异以 # NPU 注释标记:
defaults:
- ../config/traj_envs@_here_
hydra:
run:
dir: .
output_subdir: null
exp_name: "agentic_frozen_lake_npu"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: ./output/render
system_envs:
USE_MODELSCOPE: '1'
HCCL_CONNECT_TIMEOUT: "3600"
HCCL_DETERMINISTIC: "false"
HCCL_OP_EXPANSION_MODE: "AIV"
HCCL_NPU_SOCKET_PORT_RANGE: "auto"
VLLM_ASCEND_ENABLE_NZ: "0"
NPU_MEMORY_FRACTION: "0.96"
CPU_AFFINITY_CONF: "2"
OMP_NUM_THREADS: "1"
VLLM_USE_V1: "1"
checkpoint_config:
type: file_system
output_dir: /data/models/${exp_name}
num_gpus_per_node: 8
max_steps: 1024
save_steps: 10000
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false
rollout_batch_size: 1024
val_batch_size: 1024
sequence_length: 8192
advantage_clip: 0.2
ppo_epochs: 1
adv_estimator: "grpo"
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0
max_grad_norm: 1.0
pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct
actor_train:
model_args:
attn_implementation: fa2 # NPU: 通过 transformers 使用 fa2,不能使用 flash_attn
disable_gradient_checkpointing: false
dtype: bf16
model_type: ~
training_args:
learning_rate: 1.0e-6
weight_decay: 0
per_device_train_batch_size: 2
gradient_accumulation_steps: 64
warmup_steps: 10
lr_scheduler_type: cosine
data_args:
template: qwen2_5
strategy_args:
strategy_name: fsdp2_train
strategy_config:
fsdp_size: 4
param_dtype: bf16
reduce_dtype: bf16
reshard_after_forward: true
offload_policy: false # NPU: 必须使用 FSDP2,不能用 megatron_train
device_mapping: list(range(0,4)) # NPU: 训练使用 NPU 0-3
infer_batch_size: 2
actor_infer:
model_args:
disable_gradient_checkpointing: true
dtype: bf16
generating_args:
max_new_tokens: 128
top_p: 0.99
top_k: 100
num_beams: 1
temperature: 0.99
num_return_sequences: 1
data_args:
template: qwen2_5
strategy_args:
strategy_name: vllm
strategy_config:
gpu_memory_utilization: 0.8
block_size: 16
load_format: auto
device_mapping: list(range(4,8)) # NPU: 推理使用 NPU 4-7(与训练分离)
infer_batch_size: 2
reference:
model_args:
attn_implementation: fa2
disable_gradient_checkpointing: true
dtype: bf16
model_type: ~
data_args:
template: qwen2_5
strategy_args:
strategy_name: fsdp2_infer
strategy_config:
fsdp_size: 4
param_dtype: bf16
reduce_dtype: bf16
reshard_after_forward: true
offload_policy: false
device_mapping: list(range(4,8)) # NPU: 与 actor_infer 共享推理卡
infer_batch_size: 2
reward_normalization:
grouping: traj_group_id
method: mean_std
train_env_manager:
max_env_num_per_worker: 16
num_env_groups: 128
group_size: 8
tags: [FrozenLake]
num_groups_partition: [128]
val_env_manager:
max_env_num_per_worker: 32
num_env_groups: 1024
group_size: 1
tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake]
num_groups_partition: [256, 256, 256, 256]
max_tokens_per_step: 64
custom_envs:
SimpleSokoban:
${custom_env.SimpleSokoban}
LargerSokoban:
${custom_env.LargerSokoban}
SokobanDifferentGridVocab:
${custom_env.SokobanDifferentGridVocab}
FrozenLake:
${custom_env.FrozenLake}
FrozenLakeThink:
${custom_env.FrozenLakeThink}
步骤 4:启动训练
cd /workspace/ROLL
export PYTHONPATH="/workspace/ROLL:$PYTHONPATH"
python examples/start_agentic_pipeline.py \
--config_path agentic_frozen_lake_npu \
--config_name agentic_frozen_lake_npu
样例 2:单机 RLVR 流水线(Qwen3-30B-A3B)
本样例使用仓库中的 examples/ascend_examples/qwen3_30b_rlvr_fsdp2.yaml 配置在昇腾 NPU 上运行 RLVR 流水线。
关键配置
system_envs:
USE_MODELSCOPE: '1'
HCCL_CONNECT_TIMEOUT: "3600"
HCCL_DETERMINISTIC: "false"
HCCL_OP_EXPANSION_MODE: "AIV"
HCCL_NPU_SOCKET_PORT_RANGE: "auto"
VLLM_ASCEND_ENABLE_NZ: "0"
NPU_MEMORY_FRACTION: "0.96"
CPU_AFFINITY_CONF: "2"
OMP_NUM_THREADS: "1"
VLLM_USE_V1: "1"
PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
rollout_batch_size: 32
prompt_length: 2048
response_length: 4096
num_return_sequences_in_group: 8
pretrain: Qwen/Qwen3-30B-A3B
reward_pretrain: Qwen/Qwen3-30B-A3B
actor_train:
model_args:
disable_gradient_checkpointing: false
dtype: bf16
model_type: ~
training_args:
learning_rate: 1.0e-6
weight_decay: 0
per_device_train_batch_size: 2
gradient_accumulation_steps: 8
warmup_steps: 20
data_args:
template: qwen2_5
file_name:
- data/math_deepmath_deal.jsonl
domain_interleave_probs:
math_rule: 1
dataset_dir: data
messages: messages
interleave_probs: "1.0"
preprocessing_num_workers: 16
strategy_args:
strategy_name: fsdp2_train
strategy_config:
fsdp_size: 16
param_dtype: bf16
reduce_dtype: bf16
offload_policy: true
apply_expert_patch: true # NPU: MoE 模型必须启用
apply_tiled_mlp: true # NPU: TiledMLP 降低显存
tiled_num_shards: 8
reshard_after_forward: true
wrap_policy: # NPU: MoE 专用 wrap policy
wrap_embeddings: true
wrap_lm_output: true
moe_experts:
- Qwen3MoeMLP
transformer_layer_cls_to_wrap:
- Qwen3MoeAttention
- Qwen3MoeSparseMoeBlock
use_remove_padding: true
device_mapping: list(range(0,16)) # NPU: 训练使用 NPU 0-15
infer_batch_size: 2
actor_infer:
model_args:
disable_gradient_checkpointing: true
dtype: bf16
generating_args:
max_new_tokens: ${response_length}
top_p: 0.99
top_k: 100
num_beams: 1
temperature: 0.99
num_return_sequences: ${num_return_sequences_in_group}
data_args:
template: qwen2_5
strategy_args:
strategy_name: vllm
strategy_config:
gpu_memory_utilization: 0.8
block_size: 16
max_model_len: 6144
tensor_parallel_size: 2
enforce_eager: true
load_format: dummy
device_mapping: list(range(0,16)) # NPU: 推理与训练共享 NPU
infer_batch_size: 1
reference:
model_args:
disable_gradient_checkpointing: true
dtype: bf16
model_type: ~
data_args:
template: qwen2_5
strategy_args:
strategy_name: fsdp2_infer
strategy_config:
fsdp_size: 16
param_dtype: bf16
reduce_dtype: bf16
apply_tiled_mlp: true
tiled_num_shards: 8
reshard_after_forward: true
offload_policy: true
device_mapping: list(range(0,16)) # NPU: Reference 与训练共享 NPU
infer_batch_size: 2
rewards:
math_rule:
worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker
model_args:
model_name_or_path: ${reward_pretrain}
data_args:
template: qwen2_5
tag_included: [deepmath_103k, aime]
world_size: 8
infer_batch_size: 1
启动
cd /workspace/ROLL
export PYTHONPATH="/workspace/ROLL:$PYTHONPATH"
python examples/start_rlvr_pipeline.py \
--config_path ascend_examples \
--config_name qwen3_30b_rlvr_fsdp2
样例 3:多机分布式训练
本样例展示如何在多个昇腾 NPU 节点上运行 ROLL。ROLL 支持两种多机启动方式:
- 方式 A(推荐): 通过环境变量自动启动——在每个节点上设置
RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT,ROLL 自动启动和管理 Ray 集群。 - 方式 B: 手动 Ray 集群——在运行 ROLL 之前在每个节点上手动启动 Ray。
架构概览
┌──────────────────────────────────────────────────────┐
│ 主节点 (RANK=0) │
│ ┌────────────────────────────────────────────────┐ │
│ │ Docker 容器 (--net=host) │ │
│ │ ├─ Ray Head (端口 6379) │ │
│ │ ├─ Ray Dashboard (端口 8265) │ │
│ │ └─ 训练驱动 (python start_xxx.py) │ │
│ └────────────────────────────────────────────────┘ │
└──────────────────────┬───────────────────────────────┘
│ HCCL (tcp)
┌─────────────┼─────────────┐
▼ ▼
┌─────────────────────┐ ┌─────────────────────┐
│ 工作节点 1 │ │ 工作节点 2 │
│ (RANK=1) │ │ (RANK=2) │
│ ┌─────────────────┐ │ │ ┌─────────────────┐ │
│ │ Docker 容器 │ │ │ │ Docker 容器 │ │
│ │ Ray Worker │ │ │ │ Ray Worker │ │
│ │ ray start │ │ │ │ ray start │ │
│ │ --address=... │ │ │ │ --address=... │ │
│ └─────────────────┘ │ │ └─────────────────┘ │
└─────────────────────┘ └─────────────────────┘
多机前置条件
- 所有节点必须在同一二层网络内。
- 主节点的
MASTER_PORT(默认 6379)和DASHBOARD_PORT(默认 8265)必须能被所有工作节点访问(关闭防火墙或开放这些端口)。 - 所有节点需要挂载共享存储(NFS 等)到相同路径,用于模型权重、数据和断点。
- 所有节点必须使用相同的 Docker 镜像和 CANN 版本。
网络接口识别
启动前,确认每个节点上 HCCL 使用的网卡名称:
# 列出可用网卡
ip addr
# 或用 NPU 工具查看 HCCL 接口
for i in $(seq 0 7); do hccn_tool -i $i -ip -g; done
# NPU 设备 IP 通常在高速互联网络上(如 192.168.x.x)。
# 使用对应的以太网接口名称(如 enp194s0f0, eth0)作为 HCCL_SOCKET_IFNAME。
步骤 1:在所有节点启动容器
在每个节点上,使用 --net=host 启动 Docker 容器并挂载共享存储:
docker run -dit \
--name roll_npu_multi \
--ulimit nofile=65536:65536 \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci6 \
--device /dev/davinci7 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/Ascend/add-ons:/usr/local/Ascend/add-ons \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /shared/storage:/data \
--ipc=host \
--net=host \
roll:ascend-a3 \
/bin/bash
重要:
-v /shared/storage:/data将共享存储挂载到容器内,用于模型权重、训练数据和断点。此目录必须在所有节点上以相同路径访问。可使用 NFS、HDFS 或其他共享文件系统方案。