ROLL

在昇腾 NPU 上安装 ROLL,并用 FrozenLake agentic 强化学习跑通一次完整的训练闭环。

前置条件

  • 硬件:Atlas 900 A2 PODc / Ascend 910B 训练系列,单卡。

  • 软件:已装好 CANN(toolkit 与驱动),并能 source set_env.sh,Python 版本不低于 3.10。参考快速安装昇腾环境。

本文档示例在 Python 3.12、CANN 9.1.0 环境下验证通过。

本文档配套镜像:swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12。

加载 CANN 环境

source /usr/local/Ascend/ascend-toolkit/set_env.sh

安装 ROLL

使用源码安装 ROLL :以下命令均在打开终端时所在的工作目录执行,ROLL 将克隆到该目录下。

git clone --branch <ref> https://github.com/alibaba/ROLL.git
echo "ROLL $(git -C ROLL describe --tags --exact-match HEAD)"
pip install -e ROLL

Note

<ref> 为最新正式 release 的 tag

输出结果如下:

ROLL xxx

Note

输出中的 xxx 为实际安装到的 release 版本号

运行示例:FrozenLake agentic 强化学习

FrozenLake 是 ROLL 官方快速入门的示例:Qwen2.5-0.5B-Instruct 作为策略模型,在 4×4 冰面网格中逐轮输出移动方向,绕开冰洞到达终点,环境按结果返回奖励。

安装 vLLM 与 triton。 安装 vllm-ascend 时,配套的 torch、torch_npu、torchvision 和 torchaudio 会作为依赖自动安装,无需单独安装 PyTorch 软件栈。

pip install --index-url https://repo.huaweicloud.com/repository/pypi/simple vllm==0.23.0
pip install --index-url https://repo.huaweicloud.com/repository/pypi/simple --extra-index-url https://repo.huaweicloud.com/ascend/repos/pypi vllm-ascend==0.23.0
pip install --index-url https://repo.huaweicloud.com/repository/pypi/simple --extra-index-url https://repo.huaweicloud.com/ascend/repos/pypi triton-ascend==3.2.2

核对安装的 vLLM 与 triton 版本,下面代码用python执行:

from importlib.metadata import version

expected = {
    "vllm": "0.23.0",
    "vllm-ascend": "0.23.0",
    "triton-ascend": "3.2.2",
}
for package, required in expected.items():
    actual = version(package)
    assert actual == required, f"{package}: expected {required}, got {actual}"
    print(f"{package} {actual}")

输出结果如下:

vllm 0.23.0
vllm-ascend 0.23.0
triton-ascend 3.2.2

安装示例依赖。 按昇腾镜像的依赖清单安装 agentic 示例所需组件:

grep -v '^gem-llm' ROLL/requirements_common.txt | sed 's#^\./mcore_adapter$#ROLL/mcore_adapter#' > ROLL/requirements_npu.txt
sed -i 's/^decord /decord2 /' ROLL/requirements_vision.txt
pip install -r ROLL/requirements_npu.txt
pip install --ignore-requires-python gem-llm==0.0.4
pip install "numpy==1.26.4"
pip install "transformers==4.57.6" "tensorboard==2.20.0" "antlr4-python3-runtime==4.9.3"
rm ROLL/requirements_npu.txt

核对示例依赖的版本,下面代码用python执行:

from importlib.metadata import version

expected = {
    "gem-llm": "0.0.4",
    "numpy": "1.26.4",
    "transformers": "4.57.6",
    "tensorboard": "2.20.0",
    "antlr4-python3-runtime": "4.9.3",
}
for package, required in expected.items():
    actual = version(package)
    assert actual == required, f"{package}: expected {required}, got {actual}"
    print(f"{package} {actual}")

输出结果如下:

gem-llm 0.0.4
numpy 1.26.4
transformers 4.57.6
tensorboard 2.20.0
antlr4-python3-runtime 4.9.3

写入示例配置。 配置基于官方 agentic demo,修改参数完成 NPU 适配, 将配置文件保存到 ROLL/examples/agentic_frozen_lake_npu/quick_start_npu.yaml。用 Python 执行下面的代码:

from pathlib import Path

config = """
defaults:
  - ../config/traj_envs@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "roll-quick-start-npu"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: ./output/render
system_envs:
  USE_MODELSCOPE: '1'
  # RL 权重刷新场景需禁用 FRACTAL_NZ。
  VLLM_ASCEND_ENABLE_NZ: '0'
  # 允许同卡多进程由 HCCL 自动分配 device 侧端口。
  HCCL_NPU_SOCKET_PORT_RANGE: auto

track_with: tensorboard
tracker_kwargs:
  log_dir: ./output/tensorboard

num_gpus_per_node: 1

# 禁用 TransferQueue 远程传输,数据直接走 Ray 传递。
transfer_backend:
  backend_name: null

# 单卡快速跑通:只训练 1 步,批量收缩。
max_steps: 1
save_steps: 1000
logging_steps: 1
eval_steps: 1000
resume_from_checkpoint: false

rollout_batch_size: 8
val_batch_size: 2
sequence_length: 2048
max_tokens_per_step: 128

ppo_epochs: 1
adv_estimator: "grpo"
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0

pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct

actor_train:
  model_args:
    # NPU 通过 transformers 使用 fa2,不能使用 flash_attn 包。
    attn_implementation: fa2
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 1.0e-6
    per_device_train_batch_size: 1
    gradient_accumulation_steps: 2
  data_args:
    template: qwen2_5
  strategy_args:
    # NPU 不支持 Megatron,训练策略使用 FSDP2。
    strategy_name: fsdp2_train
    strategy_config:
      fsdp_size: 1
      param_dtype: bf16
      reduce_dtype: bf16
      reshard_after_forward: true
      offload_policy: false
      use_batched_model_update: false
  device_mapping: list(range(0,1))
  infer_batch_size: 1

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: 128
    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(0,1))
  infer_batch_size: 1

reference:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: true
    dtype: bf16
    model_type: ~
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: hf_infer
    strategy_config: ~
  device_mapping: list(range(0,1))
  infer_batch_size: 1

train_env_manager:
  max_env_num_per_worker: 8
  num_env_groups: 4
  group_size: 2
  tags: [FrozenLake]
  num_groups_partition: [4]

val_env_manager:
  max_env_num_per_worker: 2
  num_env_groups: 2
  group_size: 1
  tags: [FrozenLake]
  num_groups_partition: [2]

custom_envs:
  FrozenLake:
    ${custom_env.FrozenLake}
"""

path = Path("ROLL/examples/agentic_frozen_lake_npu/quick_start_npu.yaml")
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(config.lstrip("\n"), encoding="utf-8")
print("config written", path)

输出结果如下:

config written ROLL/examples/agentic_frozen_lake_npu/quick_start_npu.yaml

启动训练。 在当前工作目录运行 agentic pipeline 入口脚本,ROLL 自动拉起 Ray 集群。

python ROLL/examples/start_agentic_pipeline.py --config_path agentic_frozen_lake_npu --config_name quick_start_npu

查看训练产物。 训练完成后,TensorBoard 指标保存在当前工作目录的 output/tensorboard/roll-quick-start-npu 下。用 Python 执行下面的代码,查看训练结果指标:

from pathlib import Path

from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

tensorboard_dir = Path("output/tensorboard/roll-quick-start-npu")
event_files = sorted(tensorboard_dir.glob("*/events.out.tfevents.*"))
assert event_files, f"no tensorboard event file under {tensorboard_dir}"
metrics = EventAccumulator(str(event_files[-1]))
metrics.Reload()

print("训练产物已生成")
print(f"训练指标路径:{tensorboard_dir}")
print(f"训练步数:{len(metrics.Scalars('time/step_total'))}")
print(f"训练样本数:{int(metrics.Scalars('system/samples')[-1].value)}")
print(f"FrozenLake 平均得分:{metrics.Scalars('critic/score/mean')[-1].value:.3f}")

输出结果如下:

训练产物已生成
训练指标路径:output/tensorboard/roll-quick-start-npu
训练步数:1
训练样本数:8
FrozenLake 平均得分:xxx

Note

输出中的 xxx 为本次训练实际得到的平均得分,随训练随机性变化

外部链接