RLinf RECAP (pi0.5) 昇腾910B适配简记

发表于 2026-06-30  278 次阅读

文章目录

环境概览

组件版本/说明
硬件8 × Ascend 910B (aarch64)
基础镜像`rlinf/rlinf:agentic-rlinf0.2-libero-cann9.0`
Python3.11（自建 .venv）
PyTorch2.6.0+cpu + torch-npu 2.6.0
CANN9.0.0
transformers4.57.x
Ray2.x（分布式调度）
数据集RECAP-Libero10-Task0-48succ-Data
分布式8 卡 FSDP（hccl 通信后端）

部署步骤

1. 拉取 Docker 镜像

RLinf 官方提供了预置 CANN 9.0 环境的 Docker 镜像，镜像内包含 PyTorch 2.6、torch-npu、Ray 等核心依赖：

docker pull rlinf/rlinf:agentic-rlinf0.2-libero-cann9.0

> 注意：该镜像带 CANN 9.0 驱动栈，需确保宿主机已安装对应版本的 NPU 驱动和固件。

2. 启动容器

docker run -itd \
  --name rlinf-recap \
  --network host \
  --ipc=host \
  --privileged \
  --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/hisi_hdc \
  -v /data:/data \
  rlinf/rlinf:agentic-rlinf0.2-libero-cann9.0 \
  /bin/bash

关键参数说明：

  • --network host + --ipc=host：Ray 分布式通信必需
  • --privileged + 设备映射：NPU 设备直通
  • -v /data:/data：挂载数据目录，存放数据集和模型权重

3. 硬件可用性检查

进入容器后验证 NPU 设备：

docker exec -it rlinf-recap bash
npu-smi info

预期输出 8 卡均为 OK 状态：

NPU-ID  Name   Health  Power(Tot)  Temp(C)
0       910B   OK      65.0W       42
1       910B   OK      65.0W       41
...（8 卡均 OK）

4. 创建虚拟环境

RLinf 使用自建 .venv 而非镜像内置 conda 环境，通过 install.sh 安装依赖：

cd /RLinf
bash install.sh embodied --model openpi --env libero --platform ascend

该命令会：

  • 创建 .venv（Python 3.11）
  • 安装 torch 2.6.0+cpu + torch-npu 2.6.0
  • 安装 transformers、diffusers、accelerate、lerobot 等依赖
  • 配置 hccl 分布式通信后端

验证环境：

source .venv/bin/activate
python -c "import torch; import torch_npu; print('NPU count:', torch.npu.device_count())"
python -c "import ray; print('Ray version:', ray.__version__)"

5. 准备数据集

从 HuggingFace 下载 RECAP-Libero10 数据集：

pip install huggingface-hub

# 下载数据集到 /data 目录
huggingface-cli download \
  --repo-type dataset RLinf/RECAP-Libero10-Task0-48succ-Data \
  --local-dir /data/RECAP-Libero10-Task0-48succ-Data

数据集目录结构：

/data/RECAP-Libero10-Task0-48succ-Data/
├── libero10_task0_sft/       # 48 条成功轨迹（8005 帧），Step 2 训练集
│   ├── data/                  # 图像帧 + 动作数据
│   ├── meta/                  # episodes 元信息
│   └── info.json              # 数据集版本信息
├── libero10_task0_train/      # 4096 条成败混合 rollout
├── libero10_task0_eval/       # 验证集

6. 系统依赖补充

RLinf 推理/评估阶段可能需要额外依赖：

apt update && apt install -y \
  libavcodec-dev libavformat-dev libavutil-dev \
  libswscale-dev libswresample-dev ffmpeg

7. 关键环境变量

训练前设置以下环境变量：

# NPU 可见设备（8 卡全用）
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

# hccl 通信超时（8 卡建连需要更长时间）
export HCCL_CONNECT_TIMEOUT=600

# 关闭信号处理（避免 Ray fork 时信号冲突）
export TORCHELASTIC_ENABLE_FILE_TIMER=0

# torch-npu 库路径（确保 hccl so 可找到）
export LD_LIBRARY_PATH=/RLinf/.venv/lib/python3.11/site-packages/torch/lib:\
/RLinf/.venv/lib/python3.11/site-packages/torch_npu/lib:$LD_LIBRARY_PATH

# tf32 关闭（NPU 的 tf32 实现精度与 CUDA 不同）
export ALLOW_TF32=0

一、背景

RLinf 是一个基于 Ray 的机器人强化学习框架。其核心算法 RECAP 通过四阶段流水线实现 classifier-free guidance (CFG) 策略优化：先计算轨迹回报（Step 1），再训练价值模型（Step 2），然后计算优势标签（Step 3），最后 CFG 训练策略模型（Step 4）。整个过程涉及 VLM backbone（SigLIP + Gemma3）、8 卡 FSDP 分布式训练，目标是把这套 pipeline 完整跑在华为昇腾 910B 上。

数据集路径：

  • libero10_task0_sft：48 条成功轨迹（8005 帧），Step 2 训练集
  • libero10_task0_train：4096 条成败混合的 rollout，Step 2 eval + Step 1 输入
  • libero10_task0_eval：验证集

二、Step 1 — Compute Returns（通过）

纯 CPU 计算，无 NPU 依赖。对 8005 帧 SFT 数据和 4096 条 rollout 轨迹分别计算折扣回报（G_t = r_t + γ·G_{t+1}，γ=1.0），生成 meta/returns.parquet。全程顺畅。

python examples/recap/process/compute_returns.py --config-name=compute_returns

三、Step 2 — Value Model SFT（8 卡 FSDP 训练）

这是整个适配过程中踩坑最密集的阶段。

3.1 transfer_to_npu 的 Segfault

torch_npu.contrib.transfer_to_npu 是 torch-npu 的 monkey-patch 模块，作用是把所有 torch.cuda.* 调用自动替换为 torch.npu.*。但它在遍历 torch 模块树时触发了 transformers 的 lazy import→image_transforms→tensorflow→pywrap_tensorflow.py→self_check.py:63，TF 的 C 扩展在 NumPy 2.x 下直接 segfault。

修复：注释掉 train_value.py 里的 from torch_npu.contrib import transfer_to_npu，然后 pip uninstall tensorflow。

但实际上 Ray worker 文件 rlinf/workers/sft/fsdp_value_sft_worker.py 也 import 了它，所以每个 NPU worker fork 时还是会触发。关键是要确保环境中没有 tensorflow。

3.2 LeRobot Dataset 离线加载

Step 2 的 config 里指定了 eval_data_paths 指向本地 libero10_task0_train。LeRobot 的 LeRobotDataset.__init__ 会调用 get_safe_version()，该方法内部 get_repo_versions() 去 hf-mirror 查数据集版本信息，但这个本地数据集不在 HuggingFace 上——404。

修复：直接 patch LeRobot 源码 lerobot/common/datasets/utils.py：

# get_repo_versions — 捕捉异常返回空列表
def get_repo_versions(repo_id: str):
    try:
        return _get_repo_versions_impl(repo_id)
    except Exception:
        return []

# get_safe_version — hub_versions 为空时直接用本地版本
def get_safe_version(repo_id, version):
    hub_versions = get_repo_versions(repo_id)
    if not hub_versions:
        return f"v{target_version}"  # 用本地 info.json 的 _version_
    ...

.pth 文件和入口脚本 monkey-patch 都不生效（Ray workers 独立 fork，不继承父进程 monkey-patch），必须改源文件。

3.3 Ray Timer Crash

训练跑 4 步没问题，第 5 步触发 save+eval 时崩：

ValueError: Timer 'run_eval' has not been recorded.

根因：transfer_to_npu monkey-patch 了 torch.distributed.init_process_group（替换 backend 为 hccl），干扰了 Ray 内部 @Worker.timer 装饰器的计时机制。

修复：把 pop_execution_time 的 crash 改成 no-op：

# rlinf/scheduler/worker/worker.py:1295
if tag not in self._timer_metrics:
    return 0.0  # patched: Ray timer broken by transfer_to_npu

3.4 训练指标

8 卡 FSDP，micro_batch=16×8=128 global_batch，8005 样本，约 62 步/epoch。

前 4 步训练指标（模型尚未收敛，正常）：

  • loss ~6.0，MAE ~0.43
  • predicted_value_std ~0.01（模型输出近乎常数，预期中）
  • value_spearman ~-0.3 逐步改善（step 4 到 -0.02）

四、已完成修改清单

文件修改原因
`examples/recap/value/train_value.py:26`import transfer_to_npu（保留）NPU cuda→npu 映射
`rlinf/workers/sft/fsdp_value_sft_worker.py:30`import transfer_to_npu（保留）Ray worker 需要
`rlinf/scheduler/worker/worker.py:1295`ValueError → return 0.0Ray timer crash
`lerobot/utils.py:get_repo_versions`try/except 包裹离线数据集 404
`lerobot/utils.py:get_safe_version`hub_versions 空时用本地版同上
`examples/recap/value/config/libero_sft_value.yaml`val_check_interval=50, save_interval=50控制验证频率

备份目录：/RLinf/backups_adapt/20260630_101234/

五、Step 3 — Compute Advantages（8 卡推理）

batch_size 调优

Value Model 推理显存只用到 7G（910B 共 61G），初始 batch_size=16 时 GPU busytime 仅 30%，数据加载（avg_fetch=1088ms）远超推理（avg_infer=456ms）。调大到 batch_size=128 后 batch 数从 64 降到 8，avg_infer=2831ms，GPU busy 提升至 19% — 数据加载仍是瓶颈但吞吐有改善。

# compute_advantages_npu.yaml
advantage:
  batch_size: 128  # 16 → 128

运行结果

指标batch_size=16batch_size=128
批次数648
总耗时~104s~122s
avg_infer/批456ms2831ms
GPU busy30%19%
samples/s3723

> 注意：batch_size 翻 8 倍后总耗时反而略增（104s→122s），瓶颈在 avg_fetch=12s 的视频解码 + 预处理，不在 NPU 推理。进一步优化需提高 dataloader workers 或预缓存预处理结果。

结果统计（与 batch_size 无关）

指标值
推理样本8005 samples, 30 episodes
V(o_t) mean-0.4851
Advantage mean-0.0033
Advantage range[-0.0435, 0.5022]
Positive rate30.0% (2402/8005)
SFT 模式全部标记为 True
输出`advantages.parquet`, `mixture_config.yaml`

compute_advantages.py 修复

  • 添加 hccl 后端 override（原代码只检测 nccl）
  • 添加 torch.npu.set_device() 兼容（torch.cuda.set_device 在 NPU 上会报错）

六、Step 4 — CFG Training（8 卡 FSDP 训练）

修复清单

文件修复内容
`rlinf/workers/sft/fsdp_cfg_worker.py`NPU device 初始化：`torch.cuda.set_device` → `torch.npu.set_device`
`examples/recap/cfg/config/libero_cfg_openpi_npu.yaml`新建 NPU 配置，指向 `/data/pi05_base` + 实际数据集路径
`/data/pi05_base/physical-intelligence/libero/norm_stats.json`从数据集生成 LIBERO action/state 归一化统计（mean/std/q01/q99）
NPU OOM 修复`micro_batch_size` 32→4, `global_batch_size` 512→256

norm_stats.json 生成

pi05 通用模型缺少 LIBERO 环境专属的 norm_stats.json。从数据集的 8005 样本中提取 action(7维) 和 state(8维) 的统计信息，写入 physical-intelligence/libero/norm_stats.json：

# 从 parquet 文件读取
for f in parquet_files:
    df = pd.read_parquet(f)
    actions = np.stack(df["actions"].values)  # (N, 7)
    states = np.stack(df["state"].values)      # (N, 8)
# 计算 mean/std/q01/q99 → norm_stats.json

显存与 batch size

pi05 14B 模型在 sharding_strategy: no_shard 下每卡完整加载模型。初始 micro_batch_size=32 时前向已占 56G/61G，backward 时 OOM：

RuntimeError: NPU out of memory. Tried to allocate 970.00 MiB
  (NPU 0; 60.96 GiB total; 56.37 GiB already allocated)

减小 micro_batch_size=4 后顺利运行。gradient accumulation 自动补偿（256/(4×8)=8 步累积）。

快速验证运行（15 步）

# libero_cfg_openpi_npu_validate.yaml
runner:
  max_steps: 15          # 仅 15 步验证
  save_interval: 5       # 每 5 步保存
  val_check_interval: -1 # 关闭校验（CFG worker 不支持 eval）

actor:
  micro_batch_size: 4
  global_batch_size: 256
  optim:
    lr_warmup_steps: 3
    total_training_steps: 15
指标值
速度~13-23s/it
Loss0.0703 → 0.0433（15 步稳定下降）
Grad Norm0.815 → 0.336（收敛）
Checkpointsglobal_step_5, _10, _15 ✅
显存~56G/61G（no_shard 模式）

Eval 验证

fsdp_cfg_worker.py 硬编码 self.eval_data_loader = None，未实现 eval 逻辑（不像 fsdp_sft_worker.py 会从 val_data_paths 构建 eval dataloader）。如需 validation，需在 CFG worker 中添加 eval dataloader 支持。快速验证阶段暂时关闭校验（val_check_interval: -1）。

数据集说明

数据集EpisodesFrames用途
`libero10_task0_sft`308,005🟢 训练集（Step 2/3/4）
`libero10_task0_train`4,0961,561,874🟡 备用训练/验证集
`libero10_task0_eval`6425,493🔵 纯验证集

已知问题

  • aten::_sample_dirichlet NPU 不支持，fallback CPU
  • NO_SHARD 策略 PyTorch 已标记 deprecated，后续可改 shard_grad_op
  • 速度 ~13s/it 偏慢（pi05 14B + no_shard + 非融合算子）
  • CFG worker 不支持 eval dataloader

七、上游回馈计划

适配过程中积累的改动计划回馈 RLinf 主仓库：

优先级内容说明
P0`torch.npu.set_device()` 兼容三个 worker + compute_advantages，用 try/except 包裹，不影响 CUDA
P0FSDP device 修复`cls_embedding.weight.device` → buffer device，通用修复
P1`drop_last=True` 修复eval 死锁，对 CUDA 也有益
P1HCCL backend 自动检测`compute_advantages.py` 的分布式后端
P2NPU 配置模板`compute_advantages_npu.yaml`, `libero_cfg_openpi_npu.yaml`
P2CFG eval dataloader补充 `val_data_paths` → `eval_data_loader` 构建逻辑

完整 patch: RLinf-ascend-910b.patch，使用指南: RLinf-ascend-910b-guide.md。

简记。

赏

scanz个人博客