RLinf（main分支）Ascend NPU 适配与部署简记

文章目录

> 基于 RLinf 主分支（commit a4b6abe）适配 Ascend 910B3 NPU 的完整流程记录

---

一、背景

RLinf 是一个面向机器人基础模型的强化学习训练框架，支持大规模分布式 RL 训练。框架主分支（main）相比旧分支有大幅重构，采用 Ray + FSDP 架构，并通过 Hydra 配置系统管理实验。

本文记录将 RLinf main 分支适配到 Huawei Ascend 910B3 NPU 平台的过程，涵盖环境部署、代码修改、数据集准备和完整训练流程。

---

二、环境部署

2.1 硬件环境

项目	配置
NPU	8× Ascend 910B3，64GB HBM/卡
CPU	192 核 aarch64
内存	512GB
OS	Ubuntu 22.04 aarch64
CANN	8.3.rc1

2.2 镜像构建

基础镜像使用华为云 Ascend Hub 提供的 CANN 镜像：

FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11

完整 Dockerfile 见仓库上的 `RLinf_openpi_Dockerfile`。构建时使用国内镜像加速：

docker build -f RLinf_openpi_Dockerfile --build-arg USE_MIRRORS=1 -t rlinf-ascend:latest .

关键依赖安装命令（Dockerfile 中已集成）：

bash requirements/install.sh \
    --use-mirror \
    --platform ascend \
    embodied \
    --model openpi \
    --env maniskill_libero

2.3 容器启动

docker run -itd --name rlinf_train \
    --privileged \
    --network host \
    --shm-size=32g \
    -v /data:/data \
    rlinf-ascend:latest

---

三、Ascend NPU 适配修改

主框架代码切换到 main 分支后（`git checkout a4b6abe`），在 Ascend NPU 上遇到了三类兼容性问题。

3.1 torch.cuda API 不可用

Ascend NPU 环境安装的是 `torch 2.6.0+cpu`（不带 CUDA），所有 `torch.cuda.*` 调用都会抛 `AttributeError`。

**受影响文件和修复：**

# rlinf/workers/sft/fsdp_value_sft_worker.py
# rlinf/workers/sft/fsdp_cfg_worker.py
# 删除以下两行：
- torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
- self.device = torch.cuda.current_device()

`FSDPModelManager.__init__()` 已经通过 `Worker.torch_platform.set_device()` 自动处理了设备设置，这些行是多余的。

# rlinf/data/process/distributed.py
# 增加 NPU 分支：
elif hasattr(torch, 'npu') and torch.npu.is_available():
    torch.npu.set_device(local_rank)
    device = f"npu:{local_rank}"

3.2 torchcodec 不兼容

`torchcodec 0.14.0` 依赖 CUDA 库 `libnvrtc.so.13`，在 Ascend 上无法加载。

**修复：** 卸载 torchcodec，`LeRobotDataset` 改用 pyav 后端：

dataset = LeRobotDataset(
    dataset_path,
    download_videos=False,
    video_backend="pyav",  # 代替默认的 torchcodec
)

3.3 GPU 检测和分布式后端

`run_compute_advantages.sh` 用 `nvidia-smi` 检测 GPU 数量：

# 修改前：nvidia-smi 返回 0，wc -l 返回 exit 0，NPROC=0 → 1
NPROC_PER_NODE=$(nvidia-smi -L 2>/dev/null | wc -l || echo 1)

# 修改后：先试 nvidia，再试 torch_npu
NPROC_PER_NODE=$(nvidia-smi -L 2>/dev/null | wc -l)
if [ "$NPROC_PER_NODE" -eq 0 ] 2>/dev/null; then
    NPROC_PER_NODE=$(python3 -c "import torch; import torch_npu; print(torch.npu.device_count())")
fi

分布式后端从 `nccl` 改为 `hccl`（Huawei Collective Communication Library）：

distributed:
  enabled: true
  backend: "hccl"  # 而不是 "nccl"

---

四、数据集准备

4.1 数据目录结构

/data/
├── pi05_base/                    # Pi0.5 预训练模型权重
├── siglip2-so400m-patch14-224/   # SigLIP 视觉编码器
├── gemma-3-270m/                 # Gemma3 语言模型（270M）
└── RECAP-Libero10-Task0-48succ-Data/
    ├── libero10_task0_sft/       # 训练集：48 条成功轨迹
    │   ├── meta/
    │   ├── videos/               # 视频帧数据
    │   └── ...
    └── libero10_task0_train/     # 验证集
        └── ...

4.2 数据集配置

**训练集**: `libero10_task0_sft` — 48 条成功演示轨迹，libero10 场景
**验证集**: `libero10_task0_train` — 同场景但不同 episodes

---

五、训练流程

完整的 RECAP-CFG-RL 训练分为 4 个步骤，形成一个数据标注 → 价值学习 → 策略优化的流水线。

Step 1: Compute Returns（计算折扣回报）

输入: 原始轨迹数据（含成功/失败标签）
输出: meta/returns.parquet（每条轨迹每帧的折扣回报）

使用 `run_compute_returns.sh`，核心参数：

data:
  gamma: 1.0  # 折扣因子

Step 2: Value Model SFT（价值模型训练）

输入: 原始轨迹数据 + meta/returns.parquet
输出: 训练好的价值模型 checkpoint

使用 `run_value_sft.sh`，采用 FSDP 分片策略（no_shard），0 号卡保存检查点。在 8× 910B3 上按 30000 epoch 配置训练，val_check_interval=500，save_interval=3000。

关键配置：

参数	值	说明
micro_batch_size	16	每卡微批次
global_batch_size	128	总批次（8×16）
lr	5e-5	策略学习率
value_lr	1e-4	价值头学习率
precision	bf16	混合精度训练

Step 3: Compute Advantages（计算优势值）

输入: 原始轨迹数据 + 价值模型 checkpoint
输出: meta/advantages.parquet（每条轨迹每帧的优势值）

核心计算公式：

Aₜ = normalize(rₜ:ₜ₊ₙ) + γⁿ · V(oₜ₊ₙ) - V(oₜ)

运行在 8 卡 NPU 上，每卡独立处理一个分片。

Step 4: CFG RL Training

输入: 原始轨迹数据 + meta/advantages.parquet + pi0.5 基础模型
输出: 训练好的策略模型 checkpoint

使用 CFG（Classifier-Free Guidance）进行强化学习训练。模型架构为 OpenPI Pi0.5（SigLIP + Gemma3 270M），添加了 CFG 条件控制。

训练过程中日志示例：

Global Step: 20/31 (65%)
loss=0.068, grad_norm=0.83
conditional_ratio=0.92, unconditional_ratio=0.08

TensorBoard 指标：

指标	含义
train/loss	总训练损失
train/conditional_loss	有条件（positive CFG）损失
train/unconditional_loss	无条件损失
train/grad_norm	梯度范数
train/learning_rate	学习率（cosine schedule）
train/conditional_ratio	有条件样本比例

---

六、日志与模型输出

6.1 日志路径

logs/
├── value_sft/                    # Step 2 输出
│   └── recap_value_model_sft-{timestamp}/
│       ├── run_value_sft.log
│       ├── tensorboard/
│       └── value_sft/
│           └── checkpoints/
│               ├── global_step_10/
│               ├── global_step_20/
│               └── ...
└── cfg_rl/                       # Step 4 输出
    └── cfg_rl_openpi-{timestamp}/
        ├── run_cfg_rl.log
        ├── tensorboard/
        └── cfg_sft/
            └── checkpoints/
                └── global_step_10/

6.2 检查点结构

global_step_10/
└── actor/
    ├── dcp_checkpoint/       # 分布式 FSDP checkpoint（每卡）
    │   ├── .metadata
    │   ├── rank_0.pt
    │   └── ...
    └── model_state_dict/
        └── full_weights.pt   # 合并后的单权重文件

---

七、问题排查

7.1 Segfault 于 import transformers

**现象**: 加载 transformers 模型时 Segfault。

**原因**: `tensorflow.python.platform.self_check` 在继承 `torch_npu` 的进程中会 crash。卸载 tensorflow 即可。

pip uninstall -y tensorflow

7.2 torch.cuda 相关异常

**现象**: `AttributeError: module 'torch._C' has no attribute '_cuda_setDevice'`

**原因**: Ascend 上 torch 是 cpu 版本，无 CUDA API。

**修复**: 删除多余 `torch.cuda.set_device` 调用，改为 `torch.npu.set_device`（patch 中已处理）。

7.3 单卡检测

**现象**: 脚本显示 `GPUs: 1`，实际有 8 卡。

**原因**: `nvidia-smi` 不存在时 `wc -l` 返回 0 且 exit code 0。

**修复**: 脚本中的设备检测逻辑改为链式 fallback（patch 中已处理）。

7.4 advantage_tag 路径错误

**现象**: `cfg_rl_openpi.yaml` 中 `advantage_tag` 设为完整路径，代码却拼接成 `advantages_{完整路径}.parquet`。

**修复**: 注释掉 `advantage_tag` 项，代码自动使用默认的 `meta/advantages.parquet`。

---

八、总结

本次适配共修改 3 处核心代码和 4 个配置文件，涉及：

✅ Ascend NPU 设备管理与检测
✅ 分布式后端切换（nccl → hccl）
✅ 视频解码后端切换（torchcodec → pyav）
✅ 训练超参调优（适应 8×910B3 显存）

完整补丁文件见仓库 `rlinf-main-ascend-patch/` 目录。