RLinf 适配 OpenPI SFT 与 pi0.5 模型多卡训练简记

发表于 9 天前  83 次阅读

文章目录

最近使用 RLinf 框架适配了 pi0.5 模型的 SFT（Supervised Fine-Tuning）训练，在 Ascend 910B3 NPU 集群上跑通。过程中遇到几个坑，记录一下修改点和排查思路。

一、背景

RLinf 之前已经有 RECAP 算法的适配（基于 pi0 模型），这次需要新增 pi0.5 模型的 SFT 训练支持。模型权重已提前下载到 /data/pi05_base，数据集使用 RECAP 格式的 Libero10 子集。

容器基于之前 RECAP 适配的镜像（rlinf_train:cann83-py311-v3），Ascend CANN 环境和 torch_npu 已就绪。直接跑 bash examples/sft/run_vla_sft.sh libero_sft_openpi 会报两个错误，逐个排查。

二、问题排查与修复

2.1 is_lora 配置缺失

首次运行直接报 ConfigAttributeError: Key 'is_lora' is not in struct。

原因是 libero_sft_openpi.yaml 使用了内联的 model 配置（没有通过 defaults 引入 model/pi0_5 模板），内联配置里没有 is_lora 和 lora_rank 字段。而 rlinf/models/__init__.py 第 268 行直接 cfg.is_lora 访问，OmegaConf 的 struct 模式下会抛出异常。

修复：

  • YAML 配置中添加 is_lora: False 和 lora_rank: 32
  • 同时将代码中的 cfg.is_lora 改为 cfg.get("is_lora", False)，防御性处理
# examples/sft/config/libero_sft_openpi.yaml 片段
  model:
    model_type: "openpi"
    precision: null
    model_path: "/data/pi05_base"
    is_lora: False
    lora_rank: 32
    num_action_chunks: 4
    add_value_head: False
    openpi:
      config_name: "pi05_libero"
      detach_critic_input: True

# rlinf/models/__init__.py 第 268 行
-    if cfg.is_lora:
+    if cfg.get("is_lora", False):

2.2 JAX CUDA 初始化失败

修复 is_lora 后，数据加载阶段又报：

RuntimeError: Unable to initialize backend 'cuda': FAILED_PRECONDITION: No visible GPU devices.

服务器只有 Ascend NPU，没有 CUDA GPU。JAX 默认尝试初始化 CUDA 后端，但 openpi/training/data_loader.py 第 412 行仅仅调用 jax.process_count() 来判断是否多进程场景——这个调用完全可以用 CPU 后端满足。

修复：在启动脚本中添加环境变量：

# examples/sft/run_vla_sft.sh
export JAX_PLATFORMS=cpu

2.3 Ascend 自动迁移（transfer_to_npu）

在 Ascend NPU 上跑 PyTorch 代码时，torch_npu 的 transfer_to_npu 模块会自动将 torch.cuda.* 调用替换为 torch.npu.*。这是由 fsdp_sft_worker.py 顶层的 import torch_npu 触发的。日志中会看到如下提示：

The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
The backend in torch.distributed.init_process_group set to hccl now..

此功能为 Ascend CANN 内置，无需额外配置。已经包含在容器镜像的 torch_npu 包中。

三、多卡配置

默认配置只使用 1 张 NPU（actor,env,rollout: 0-0），需要改为 8 卡全量训练：

# examples/sft/config/libero_sft_openpi.yaml
cluster:
  num_nodes: 1
  component_placement:
    actor,env,rollout: 0-7    # 使用全部 8 张 Ascend 910B3

RLinf 的 FlexiblePlacementStrategy 会根据 component_placement 自动分配 GPU 范围。FSDP 的 full_shard 策略会自动跨 8 卡分片模型参数。单卡验证时每步约 62 秒（pi0.5 约 14GB 参数量），8 卡后每步内存占用降低、吞吐量应有显著提升。

参数单卡（1×910B3）8卡（8×910B3）
component_placementactor,env,rollout: 0-0actor,env,rollout: 0-7
micro_batch_size44
global_batch_size128128
FSDP shardingNO_SHARD（world=1）FULL_SHARD
每步时间（实测）~62s待验证

四、验证结果

修复后在单张 NPU 上验证通过，训练日志输出：

Global Step: 1/30000 [01:01<514:18:22, 61.72s/it
time/step=61.7, train/grad_norm=0.533
train/learning_rate=2.5e-8, train/loss=0.0633

NPU 0 资源使用情况：

  • AICore 利用率：33%
  • HBM 占用：65494/65536 MB（模型加载）
  • 功耗：190.6W
  • NPU 进程：rayFSDPVlaSft 占用 62GB

五、Patch 使用

所有修改汇总为一份 patch 文件，可在其他节点直接应用：

# 在 RLinf 仓库根目录应用
cd /workspace/RLinf
git apply rlinf-openpi-sft-pi05-8card.patch

Patch 包含：

  • examples/sft/config/libero_sft_openpi.yaml：添加 is_lora + 8卡 placement
  • examples/sft/run_vla_sft.sh：添加 JAX_PLATFORMS=cpu
  • rlinf/models/__init__.py：cfg.is_lora → cfg.get 防御

应用后启动训练：

cd /workspace/RLinf
source .venv/bin/activate
bash examples/sft/run_vla_sft.sh libero_sft_openpi

简记。

赏

scanz个人博客