在百舸上使用 RLinf加速镜像训练 OpenPI
概述
RLinf OpenPI Torch 加速镜像面向具身策略模型监督微调(SFT)场景,在社区官方镜像的OpenPI Torch版本对齐JAX版精度的基础之上,进行了性能优化,并将官方镜像基础环境进行了 aihc.lgn7ib 实例架构适配。
通过百舸平台,您可以直接选择已适配的训练镜像、GPU 资源规格、数据集和模型权重,使用通用环境变量启动单机或多机分布式训练任务,无需手动拉起 Docker 容器或自行维护底层运行环境。
适用场景
- 使用 LIBERO 数据集训练 OpenPI π0.5 策略模型。
- 基于
pi05_base基座权重进行 Full SFT 或 LoRA SFT。 - 在百舸平台上验证单机 / 多机 GPU 训练性能。
- 需要 torch 2.13 + CUDA 13.2 + cuBLAS 13.5 加速运行时的场景。
环境与资源要求
| 项目 | 推荐配置 |
|---|---|
| GPU 规格 | aihc.lgn7ib ,单卡约 96-97 GiB |
| 单机训练 | 单机 8 卡 |
| 多机训练 | 双机 16 卡,可按需扩展 |
| 工作目录 | /root/workspace/RLinf |
| 启动脚本 | /root/workspace/run_libero_openpi_rlinf.sh |
| 训练入口 | /root/workspace/RLinf/examples/sft/train_vla_sft.py |
| 默认配置 | libero_sft_openpi_rlinf |
| 数据集目录 | /root/workspace/asset/datasets/physical-intelligence/libero |
| 模型权重目录 | /root/workspace/asset/models/pi05_base_new |
| Python 环境 | /opt/venv/openpi |
说明:本文中的路径均以百舸训练任务内的 /root/workspace 为例。实际使用时,可根据任务挂载目录调整。
加速镜像软件栈
本镜像与官方 agentic-rlinf0.3-maniskill_libero(torch 2.11.0+cu130)的核心差异在 openpi 运行环境,实测版本如下:
| 组件 | 版本 |
|---|---|
| Python | 3.12.13 |
| torch | 2.13.0+cu132 |
| torch CUDA runtime | 13.2 |
| nvidia-cublas | 13.5.1.27 |
| nvidia-cudnn | 9.20.0.48(92000) |
| NCCL | 2.29.7 |
| triton | 3.7.1 |
| flash-attn | 2.8.3.post1 |
| jax / jaxlib | 0.11.0 |
| numpy | 2.5.2 |
| torchvision | 0.28.0+cu132 |
| torchcodec | 0.15.0+cu130 |
| lerobot | 0.1.0 |
| ray | 2.56.0 |
| openpi-client | 0.1.2 |
| rlinf-openpi | 0.1.1 |
| rlinf-transformer-openpi | 4.53.2 |
准备工作
建议先配置 HuggingFace 镜像源(如需从 HuggingFace 拉取资源):
1export HF_ENDPOINT=https://hf-mirror.com
2export HF_HUB_ENABLE_HF_TRANSFER=1
百舸平台已提供数据集和模型权重的 BOS 地址,建议直接从 BOS 下载或使用平台挂载,避免依赖外部网络。
步骤一:连接并预检资源
连接开发机后,先确认 GPU、显存、数据集、权重、镜像环境是否就绪:
1nvidia-smi
2nvidia-smi --query-gpu=index,name,memory.total,memory.used --format=csv
3free -h
4ls -ld /root/workspace/RLinf
5ls /root/workspace/run_libero_openpi_rlinf.sh
6ls -ld /root/workspace/asset/datasets/physical-intelligence/libero
7ls -l /root/workspace/asset/models/pi05_base_new/model.safetensors
确认 openpi 运行环境:
1/opt/venv/openpi/bin/python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_arch_list())"
预期输出 torch 2.13.0+cu132、13.2。
步骤二:准备数据集
OpenPI SFT 要求数据集为 LeRobot format,目录必须包含 meta/info.json。
1mkdir -p /root/workspace/asset/datasets/physical-intelligence/libero
2bcecmd bos cp -r bos:/aihc-rdw-bj/huggingface.co/datasets/physical-intelligence/libero/ \
3 /root/workspace/asset/datasets/physical-intelligence/libero/ --yes
验证数据集结构:
1python3 - <<'PY'
2from pathlib import Path
3root = Path('/root/workspace/asset/datasets/physical-intelligence/libero')
4assert (root / 'meta/info.json').exists(), 'meta/info.json 不存在'
5print('parquet:', len(list(root.glob('data/**/*.parquet'))))
6print('videos:', len(list(root.glob('videos/**/*'))))
7PY
步骤三:准备模型权重与转换
训练时 BASE_MODEL_PATH 必须指向包含 model.safetensors 的 PyTorch 原生(new 格式)checkpoint 目录(本镜像默认使用 pi05_base_new)。BOS 上没有 new 格式,只有 old 格式的 pi05_base,需要本地执行 old→new 转换。
涉及三个目录:
| 目录 | 格式 | 用途 |
|---|---|---|
pi05_base |
old 格式 | 转换输入,仅需 model.safetensors + config.json 两个文件 |
pi05_base_new |
new 格式 | 训练的 model_path(BASE_MODEL_PATH 默认值) |
RLinf-Pi05-LIBERO-SFT |
assets | 存放 norm stats |
3.1 下载 old 格式基座与 norm stats
1#!/usr/bin/env bash
2set -eo pipefail
3
4PFS_ROOT=/root/workspace/asset
5BASE_DIR=${PFS_ROOT}/models/pi05_base # old 格式,转换的输入
6MODEL_DIR=${PFS_ROOT}/models/pi05_base_new # new 格式,训练的 model_path
7ASSETS_DIR=${PFS_ROOT}/models/rlinf/RLinf-Pi05-LIBERO-SFT
8NORM_STATS=${ASSETS_DIR}/physical-intelligence/libero/norm_stats.json
9
10mkdir -p "${BASE_DIR}" "${MODEL_DIR}" "${ASSETS_DIR}/physical-intelligence/libero"
11
12# pi0.5 基座权重,old 格式。只需两个文件
13bcecmd bos cp "bos:/aihc-models-bj/lerobot/pi05_base/model.safetensors" "${BASE_DIR}/model.safetensors"
14bcecmd bos cp "bos:/aihc-models-bj/lerobot/pi05_base/config.json" "${BASE_DIR}/config.json"
15
16# norm stats:整仓 7.5 G,只取这一个 4507 B 的文件
17bcecmd bos cp \
18 "bos:/aihc-models-bj/RLinf/RLinf-Pi05-LIBERO-SFT/physical-intelligence/libero/norm_stats.json" \
19 "${NORM_STATS}"
3.2 old → new 格式转换
1#!/usr/bin/env bash
2set -eo pipefail
3
4PFS_ROOT=/root/workspace/asset
5DATASET_DIR=${PFS_ROOT}/datasets/physical-intelligence/libero
6BASE_DIR=${PFS_ROOT}/models/pi05_base # old 格式,转换的输入
7MODEL_DIR=${PFS_ROOT}/models/pi05_base_new # new 格式,训练的 model_path
8ASSETS_DIR=${PFS_ROOT}/models/rlinf/RLinf-Pi05-LIBERO-SFT
9NORM_STATS=${ASSETS_DIR}/physical-intelligence/libero/norm_stats.json
10
11# 4) old -> new 转换(BOS 无 new 格式,必做)
12# 纯 CPU,峰值约 30 G 内存;丢弃 action-expert 的 lm_head,
13# 812 张量/3.6168B -> 667 张量/3.3534B,不改 dtype(仍 fp32)
14cd "${REPO_PATH:-/root/workspace/RLinf}"
15python -m rlinf.utils.ckpt_convertor.openpi.convert old2new \
16 --input-model "${BASE_DIR}/model.safetensors" \
17 --input-norm-stats "${NORM_STATS}" \
18 --output-model "${MODEL_DIR}" \
19 --output-norm-stats "${MODEL_DIR}/physical-intelligence/libero/norm_stats.json"
20
21# 5) 校验(只读)。期望 model.safetensors 13413809816 B,两份 norm_stats sha256 相同
22ls -l "${MODEL_DIR}/model.safetensors" "${NORM_STATS}"
23sha256sum "${NORM_STATS}" "${MODEL_DIR}/physical-intelligence/libero/norm_stats.json"
24du -sh "${DATASET_DIR}" "${BASE_DIR}" "${MODEL_DIR}"
25# 转换后 BASE_DIR(13.5 GiB)可删,训练只读 MODEL_DIR
转换说明:
- 纯 CPU 运行,峰值约 30 G 内存。
- 丢弃 action-expert 的
lm_head,张量数 812 → 667(3.6168B → 3.3534B)。 - 不改 dtype,仍为 fp32。
3.3 校验
1# 期望 model.safetensors 为 13413809816 B,两份 norm_stats sha256 相同
2ls -l "${MODEL_DIR}/model.safetensors" "${NORM_STATS}"
3sha256sum "${NORM_STATS}" "${MODEL_DIR}/physical-intelligence/libero/norm_stats.json"
4du -sh "${PFS_ROOT}/datasets/physical-intelligence/libero" "${BASE_DIR}" "${MODEL_DIR}"
转换完成后 BASE_DIR(约 13.5 GiB)可删除,训练只读 MODEL_DIR。转换产物 ${MODEL_DIR}/physical-intelligence/libero/norm_stats.json 即训练时 openpi 读取的 norm stats(与 ASSETS_PATH 指向同一目录 /mnt/pfs/models/pi05_base_new)。
步骤四:配置训练环境变量
启动脚本 /root/workspace/run_libero_openpi_rlinf.sh 已集中定义全部可覆盖变量,路径均可通过通用环境变量调整:
| 变量 | 默认值 | 含义 |
|---|---|---|
WORKSPACE_DIR |
/root/workspace/RLinf |
仓库根 |
DATA_PATH |
/root/workspace/asset/datasets/physical-intelligence/libero |
LIBERO 数据集根 |
BASE_MODEL_PATH |
/root/workspace/asset/models/pi05_base_new |
pi0.5 基座权重 |
ASSETS_PATH |
/root/workspace/asset/models/pi05_base_new |
openpi assets |
PROFILE |
opt_all |
baseline / mbs32 / opt_all |
NUM_NODES |
1 | 节点数(多机时设 >1) |
GBS |
自动 = mbs×8×NUM_NODES | 全局 batch size |
RAY_PORT |
6379 | Ray GCS 端口 |
RLINF_RENDEZVOUS_DIR |
/mnt/pfs/tmp/rlinf |
多机 rendezvous 共享目录,需要用共享存储 |
脚本内部会自动把通用变量映射回下游 config 硬依赖的旧变量名(LIBERO_LEROBOT_ROOT / PI05_BASE_PYTORCH_CKPT / OPENPI_ASSETS_DIR),训练逻辑零改动,无需手工 export。
步骤五:启动训练
单机 8 卡
1bash /root/workspace/run_libero_openpi_rlinf.sh
多机(PyTorchJob,所有 pod 同一行,靠 RANK/WORLD_SIZE 自动分工)
1bash /root/workspace/run_libero_openpi_rlinf.sh actor.micro_batch_size=32
通用变量覆盖示例
1WORKSPACE_DIR=/root/workspace/RLinf \
2DATA_PATH=/root/workspace/asset/datasets/physical-intelligence/libero \
3BASE_MODEL_PATH=/root/workspace/asset/models/pi05_base_new \
4ASSETS_PATH=/root/workspace/asset/models/pi05_base_new \
5PROFILE=opt_all \
6bash /root/workspace/run_libero_openpi_rlinf.sh
PROFILE=opt_all 会启用 FSDP 分片、torch.compile、fused AdamW、数据预取、图像管线等 12 项优化(+opt/... 组)。训练日志默认输出到 ${WORKSPACE_DIR}/logs/<timestamp>-libero_sft_openpi_rlinf/。
步骤六:监控训练
查看日志:
1tail -f /root/workspace/RLinf/logs/<timestamp>-libero_sft_openpi_rlinf/run_embodiment.log
查看 GPU:
1watch -n 2 nvidia-smi
查看 Ray 集群:
1ray status
性能参考
在 aihc.lgn7ib 实例单机 8 卡、PROFILE=opt_all、actor.global_batch_size=32 测试环境下,参考性能如下:
| 指标 | 参考值 |
|---|---|
| Step Time | 1399.2 ms |
| Samples/s | 182.99 |
以上数据为特定测试环境下的参考值,实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。
常见问题
1. 数据集路径不存在
默认数据集为 /mnt/pfs/datasets/physical-intelligence/libero。若挂载位置不同,通过 DATA_PATH 覆盖即可:
1DATA_PATH=/your/libero/dataset bash /root/workspace/run_libero_openpi_rlinf.sh
2. 模型权重缺少 model.safetensors
BASE_MODEL_PATH 必须指向包含 model.safetensors 的 PyTorch 原生 checkpoint 目录。确认:
1ls -lh /root/workspace/asset/models/pi05_base_new/model.safetensors
3. torch 2.13 首次训练崩溃(Cannot access data pointer of Tensor)
torch 2.13 的 inductor 缓存不可删除,否则首次训练在 Triton GEMM autotune 冷缓存下可能崩溃。制作镜像时已把预热缓存一并打入,请勿清空 torchinductor_cache_213。
4. 多机 worker 未接入 Ray
确认 RLINF_RENDEZVOUS_DIR 共享盘可读写,master 已写入 head IP;worker 会轮询等待 head 地址(默认超时 1800s)。可临时调大 RAY_WAIT_TIMEOUT:
1RAY_WAIT_TIMEOUT=3600 NUM_NODES=2 bash /root/workspace/run_libero_openpi_rlinf.sh
评价此篇文章
