在百舸上使用 openpi (JAX) 加速镜像训练 pi0.5 模型
一、概述
本文介绍如何在百舸平台使用 openpi（JAX 版）pi0.5 训练加速镜像，完成 pi0.5 模型在 LIBERO 数据集上的监督微调（SFT）训练。该加速镜像在保持训练效果与收敛稳定性（等价性）的前提下，通过数据供给优化与计算侧优化，大幅缩短单轮迭代耗时、提升训练吞吐与 GPU 利用率，并已打通单机 8 卡与多机（2 节点）数据并行训练全链路。
openpi 训练框架为 JAX 0.11 + flax nnx 的单控制器 SPMD 数据并行（非 PyTorch DDP / torchrun），由 jax.jit + NamedSharding 驱动多卡、XLA 自动插入集合通信，每节点启动 1 个训练进程。
二、适用场景
- 在 lgn7ib GPU 资源上训练 / 微调 pi0.5（π0.5）VLA 模型。
- 基于 LIBERO 数据集进行 pi05_libero SFT 训练。
- 需要单机 8 卡高吞吐训练，或多机 2 节点数据并行横向扩展。
- 希望开箱即用、免环境搭建的 JAX 0.11 训练环境。
三、模型简介
- pi0.5（π0.5） 是 Physical Intelligence 开源的 VLA（Vision-Language-Action）模型，openpi 为其官方实现（含 JAX / PyTorch 两条路径，本镜像为 JAX 路径）。
- 结构：PaliGemma VLM（SigLIP 视觉塔 + Gemma LLM）+ action expert，以 flow-matching 生成动作 chunk；gemma 主干 18 层以
nn.scan（→ XLA while loop）堆叠。 - 本任务
pi05_libero：在 LIBERO（40 个桌面操作任务）数据上做 SFT，动作 chunk 长度action_horizon=10。 - 官方仓库：https://github.com/Physical-Intelligence/openpi
四、环境与资源要求
| 项 | 说明 |
|---|---|
| GPU 资源 | lgn7ib，单卡约 96GB，推荐单机 8 卡；多机场景推荐 2 节点 × 8 卡 |
| 代码目录 | /root/workspace/openpi |
| 运行时环境 | 镜像自带 JAX 0.11 栈，开箱即用，无需挂载 venv / conda |
五、准备工作
本镜像的训练过程包含进入终端、配置环境变量、软链数据、按脚本启动训练、查看日志等交互式操作，建议在百舸开发机中执行：
- 在百舸平台创建 / 进入开发机，选择 lgn7ib GPU 资源规格（推荐 8 卡）与上述训练镜像。
- 打开开发机终端（或 WebIDE 终端），进入代码目录
/root/workspace/openpi。 - 按后续章节顺序准备数据集、准备模型权重、配置训练环境变量、启动训练。
- 在开发机中查看训练日志与产物。
六、准备数据集
LIBERO 数据集（lerobot 格式，repo_id=physical-intelligence/libero，1693 episode / 27.3 万帧）。
方式一：从 BOS 下载（推荐）
1mkdir -p /root/workspace/datasets/libero
2bosutil cp -r bos://aihc-rdw-bj/huggingface.co/datasets/physical-intelligence/libero/ /root/workspace/datasets/libero/方式二：从 HuggingFace 下载
1export HF_ENDPOINT=https://hf-mirror.com
2export HF_HUB_ENABLE_HF_TRANSFER=1
3export HF_HOME=/root/workspace/models/huggingface
4export HF_HUB_CACHE=$HF_HOME/hub
5export TRANSFORMERS_CACHE=$HF_HOME/transformers
6export HF_DATASETS_CACHE=$HF_HOME/datasets
7
8huggingface-cli download physical-intelligence/libero --repo-type dataset --local-dir /root/workspace/datasets/libero七、准备模型权重
pi05_base 基座权重（JAX / Orbax 格式）。
方式一：从 BOS 下载（推荐）
1mkdir -p /root/workspace/models/pi05_base/params
2bosutil cp -r bos:/aihc-models-bj/gs/openpi-assets/pi05_base/params/ /root/workspace/models/pi05_base/params/方式二：从 Google Cloud Storage 下载（需联网）
1# GCS Orbax 源：gs://openpi-assets/checkpoints/pi05_base/params八、配置训练环境变量
数据与权重不入镜像，运行时通过标准环境变量指定。各集群挂载点不同，用你的实际路径覆盖对应变量即可：
1export DATA_PATH="${DATA_PATH:-/root/workspace/datasets/libero}"
2export BASE_CHECKPOINT_PATH="${BASE_CHECKPOINT_PATH:-/root/workspace/models/pi05_base/params}"
3
4if [[ -z "${HF_HOME:-}" || "${HF_HOME}" == /app/* ]]; then
5 export HF_HOME=/root/workspace/openpi/.cache/hf
6fi
7export HF_LEROBOT_HOME="${HF_LEROBOT_HOME:-${HF_HOME}/lerobot}"
8mkdir -p "$HF_LEROBOT_HOME/physical-intelligence"
9
10ln -sfn "$DATA_PATH" "$HF_LEROBOT_HOME/physical-intelligence/libero"数据目录按 lerobot 格式组织：physical-intelligence/libero/{data,meta,videos}。norm stats 已随镜像提供于 assets/；若更换数据集需重算：
1python scripts/compute_norm_stats.py --config-name pi05_libero --num-workers 48启动训练前，请确认 GPU 空闲、数据 / 权重已就位、磁盘可用空间充足，并为本次实验使用独立输出目录，避免覆盖历史实验。
九、启动训练
9.1 关键参数说明
| 参数 | 值 | 说明 |
|---|---|---|
--batch-size |
256（默认） | 全局 batch（跨 8 卡） |
--num-workers |
8（优化）/ 2（默认） | DataLoader worker 数，本负载主瓶颈 |
--fsdp-devices |
1（默认，纯 DP） | >1 走 ZeRO-3 分片，省显存慢 6–8%；追吞吐峰值 / 多机弱扩展时使用 |
--model.max-token-len |
28（优化）/ 200（默认） | libero 实测 max=21，降到 28 可省 FFN 计算 |
--weight-loader.params-path |
${BASE_CHECKPOINT_PATH} |
base 权重路径 |
--num-train-steps |
30000 | 总步数 |
9.2 优化开关
以下开关默认关闭（等于上游行为），run_opt.sh 已默认开启全部优化：
1export OPENPI_SKIP_VIDEO_DECODE=1 # 跳过未消费的视频解码（占 71% 数据成本）
2export OPENPI_SKIP_IMAGES=right_wrist_0_rgb # 跳过恒 masked 的右腕图 token
3export OPENPI_FLASH_ATTN=1 OPENPI_FLASH_IMPL=xla # flash attention（XLA 实现）
4export OPENPI_WORKER_CPU=1 # 数据 worker 赶到 CPU，不抢 GPU
5# CLI：--num-workers 8 --model.max-token-len 28
6# config-B XLA flag（大 batch 下 +2.9%；OPENPI_STRICT_EQUIV=1 可关这组 flag）：
7export JAX_OPTIMIZATION_LEVEL=O1
8export XLA_FLAGS="--xla_gpu_triton_gemm_any=false --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_pipelined_all_gather=true"9.3 启动脚本
容器默认工作目录 /root/workspace/openpi，三个脚本均在此目录：
1# 基线（上游默认：nw2 / tok200 / 无优化）
2bash /root/workspace/openpi/run_baseline.sh
3# 优化（全优化）
4bash /root/workspace/openpi/run_opt.sh
5# 追吞吐峰值（大 batch，需 fsdp 分片腾显存）
6bash /root/workspace/openpi/run_opt.sh --fsdp-devices 8 --batch-size 2048脚本支持 EXP_NAME=xxx、RESUME=1、SAVE_INTERVAL=n 等环境变量，以及任意 train.py 参数透传。
9.4 断点续训
- 设置
RESUME=1时脚本走--resume，由 Orbax CheckpointManager 从<checkpoint-base-dir>/pi05_libero/<EXP_NAME>恢复完整train_state；未设置时走--overwrite。--resume与--overwrite互斥。 - checkpoint 路径模板：
<checkpoint-base-dir>/pi05_libero/<EXP_NAME>/<loop_index>；单个 checkpoint（params + optimizer + EMA）约 42GB，max_to_keep=1会清理同一实验中较旧的 checkpoint。 - 存盘语义：
SAVE_INTERVAL=n表示每 n 步存一次；训练最后一步必定会存一次 checkpoint。checkpoint 根目录用--checkpoint-base-dir <共享存储路径>指定，必须落共享文件系统。
1# 首次训练（存 ckpt）
2EXP_NAME=myrun SAVE_INTERVAL=1000 bash /root/workspace/openpi/run_opt.sh --num-train-steps 30000
3# 中断后续训（同名 EXP_NAME + RESUME=1）
4EXP_NAME=myrun RESUME=1 bash /root/workspace/openpi/run_opt.sh --num-train-steps 30000十、多机训练（2 节点数据并行）
openpi JAX 为 SPMD 单控制器，每节点启动 1 个进程执行 python scripts/train.py，该进程经 jax.jit + NamedSharding 驱动本节点全部 8 卡（不是每卡 1 进程）。各节点运行同一条命令，平台注入以下变量：
| 平台 env | 含义 | JAX 映射 |
|---|---|---|
MASTER_ADDR / MASTER_PORT |
coordinator(rank0) 地址:端口 | coordinator_address |
RANK |
节点 index（0=coordinator） | process_id |
WORLD_SIZE |
节点数（非卡数） | num_processes |
NPROC_PER_NODE |
每节点卡数 | 仅信息用，JAX 自动探测本地卡 |
启动脚本 run_multinode.sh 已封装上述变量读取、NCCL / IB 环境与优化配方。
容器 RDMA / RoCE 要求（关键）：多机需容器能做 RDMA 内存注册，创建容器时必须具备以下能力：
1--network host --ipc host --gpus all \
2--device /dev/infiniband --cap-add=IPC_LOCK --ulimit memlock=-1 --ulimit stack=67108864- 缺
IPC_LOCK/memlock=-1会报ibv_reg_mr_iova2 failed: Cannot allocate memory或NET/IB: No device found。 - 若网卡为 RoCE（link_layer=Ethernet），需按需设置
NCCL_IB_GID_INDEX=3。 - 集合通信依赖 IB / RoCE，
NCCL_IB_DISABLE=1强制走 socket 会导致 rendezvous 超时。
十一、性能参考
以下数据为特定测试环境下的参考值。实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。
单机 8 卡加速比（bs256 全局 batch，fsdp1）：
| 配置 | step time (s) | samples/s |
|---|---|---|
| baseline | 20.10 | 12.73 |
| 优化 | 1.57 | 163.01 |
| 加速比 | 12.8× | 12.8× |
baseline 吞吐偏低的根因是上游默认
num_workers=2导致数据供给严重不足（GPU 几乎空等）；优化后转为 compute-bound。优化收益主力来自数据供给（worker 数提升 + worker CPU 隔离 + 跳过无用视频解码），计算侧优化（跳过 masked 图 token / flash-xla / 降低 max-token-len）为锦上添花。
单机吞吐峰值：
| 配置 | samples/s |
|---|---|
| 优化 fsdp1 bs256（默认） | 163.01 |
| 优化 fsdp8 bs2048（追峰值） | 228.88 |
双机 2×8 弱扩展：
| 配置 | samples/s | 扩展比 |
|---|---|---|
| 单机 fsdp8 bs2048 | 228.88 | 1×（基准） |
| 双机 fsdp16 bs4096 | 443.41 | 1.94× / 96.9% |
等价性说明：优化版的 --model.max-token-len 28 与 OPENPI_SKIP_IMAGES=right_wrist_0_rgb 改变了输入张量形状，带来约 1e-4 的浮点漂移；两者数学上只影响被 mask 的无效 token，收敛终态一致，step0 两者仅差约 1e-4。数据顺序由 seed=42 固定、可复现。因此优化前后 loss 收敛曲线一致，训练结果等价。
十二、常见问题
Q1：多机训练 checkpoint 存盘卡在 Waiting for Save Finalize？
checkpoint 目录必须落在所有节点同一路径可见的共享文件系统上。orbax 多进程 finalize 需跨进程协调 rename，若各节点写各自本地盘，finalize 屏障永不达成。请通过 --checkpoint-base-dir <共享存储路径> 指定共享存储目录。
Q2：多机启动报 ibv_reg_mr_iova2 failed 或 NET/IB: No device found？
容器缺少 RDMA 内存注册能力，请按「十、多机训练」中的要求加上 --device /dev/infiniband --cap-add=IPC_LOCK --ulimit memlock=-1。
Q3：显存不足（OOM）？
优先使用 --fsdp-devices >1 走 ZeRO-3 分片省显存（会慢 6–8%），或适当降低 --batch-size。
Q4：离线环境下 lerobot 访问 HuggingFace Hub 报 Network unreachable？
设置离线开关并确保数据软链已建立：
1export HF_HUB_OFFLINE=1
2export TRANSFORMERS_OFFLINE=1十三、后续操作
- 训练完成后，checkpoint 落于
--checkpoint-base-dir指定的共享存储目录，可用于后续评测或继续训练。 - 如需长训，建议配置合理的
SAVE_INTERVAL并使用独立的EXP_NAME，便于断点续训与版本管理。 - 更换数据集时，请重新计算 norm stats（见第八章）。
评价此篇文章