{"componentChunkName":"component---src-templates-acg-portal-new-template-tsx","path":"/Dmvs45f6e","result":{"data":{"markdownRemark":{"html":"

在百舸上使用 openpi (JAX) 加速镜像训练 pi0.5 模型

\n

一、概述

\n

本文介绍如何在百舸平台使用 openpi（JAX 版）pi0.5 训练加速镜像，完成 pi0.5 模型在 LIBERO 数据集上的监督微调（SFT）训练。该加速镜像在保持训练效果与收敛稳定性（等价性）的前提下，通过数据供给优化与计算侧优化，大幅缩短单轮迭代耗时、提升训练吞吐与 GPU 利用率，并已打通单机 8 卡与多机（2 节点）数据并行训练全链路。

\n

openpi 训练框架为 JAX 0.11 + flax nnx 的单控制器 SPMD 数据并行（非 PyTorch DDP / torchrun），由 jax.jit + NamedSharding 驱动多卡、XLA 自动插入集合通信，每节点启动 1 个训练进程。

\n

二、适用场景

\n\n

三、模型简介

\n\n

四、环境与资源要求

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
项说明
GPU 资源lgn7ib，单卡约 96GB，推荐单机 8 卡；多机场景推荐 2 节点 × 8 卡
代码目录/root/workspace/openpi
运行时环境镜像自带 JAX 0.11 栈，开箱即用，无需挂载 venv / conda
\n

五、准备工作

\n

本镜像的训练过程包含进入终端、配置环境变量、软链数据、按脚本启动训练、查看日志等交互式操作，建议在百舸开发机中执行：

\n
    \n
  1. 在百舸平台创建 / 进入开发机，选择 lgn7ib GPU 资源规格（推荐 8 卡）与上述训练镜像。
    2. \n
  2. 打开开发机终端（或 WebIDE 终端），进入代码目录 /root/workspace/openpi。
    3. \n
  3. 按后续章节顺序准备数据集、准备模型权重、配置训练环境变量、启动训练。
    4. \n
  4. 在开发机中查看训练日志与产物。
    5. \n
\n

六、准备数据集

\n

LIBERO 数据集（lerobot 格式，repo_id=physical-intelligence/libero，1693 episode / 27.3 万帧）。

\n

方式一：从 BOS 下载（推荐）

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1mkdir -p /root/workspace/datasets/libero\n2bosutil cp -r bos://aihc-rdw-bj/huggingface.co/datasets/physical-intelligence/libero/ /root/workspace/datasets/libero/
\n
\n
\n
\n \n

方式二：从 HuggingFace 下载

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1export HF_ENDPOINT=https://hf-mirror.com\n2export HF_HUB_ENABLE_HF_TRANSFER=1\n3export HF_HOME=/root/workspace/models/huggingface\n4export HF_HUB_CACHE=$HF_HOME/hub\n5export TRANSFORMERS_CACHE=$HF_HOME/transformers\n6export HF_DATASETS_CACHE=$HF_HOME/datasets\n7\n8huggingface-cli download physical-intelligence/libero --repo-type dataset --local-dir /root/workspace/datasets/libero
\n
\n
\n
\n \n

七、准备模型权重

\n

pi05_base 基座权重（JAX / Orbax 格式）。

\n

方式一：从 BOS 下载（推荐）

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1mkdir -p /root/workspace/models/pi05_base/params\n2bosutil cp -r bos:/aihc-models-bj/gs/openpi-assets/pi05_base/params/ /root/workspace/models/pi05_base/params/
\n
\n
\n
\n \n

方式二：从 Google Cloud Storage 下载（需联网）

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1# GCS Orbax 源：gs://openpi-assets/checkpoints/pi05_base/params
\n
\n
\n
\n \n

八、配置训练环境变量

\n

数据与权重不入镜像，运行时通过标准环境变量指定。各集群挂载点不同，用你的实际路径覆盖对应变量即可：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1export DATA_PATH=\"${DATA_PATH:-/root/workspace/datasets/libero}\"\n2export BASE_CHECKPOINT_PATH=\"${BASE_CHECKPOINT_PATH:-/root/workspace/models/pi05_base/params}\"\n3\n4if [[ -z \"${HF_HOME:-}\" || \"${HF_HOME}\" == /app/* ]]; then\n5    export HF_HOME=/root/workspace/openpi/.cache/hf\n6fi\n7export HF_LEROBOT_HOME=\"${HF_LEROBOT_HOME:-${HF_HOME}/lerobot}\"\n8mkdir -p \"$HF_LEROBOT_HOME/physical-intelligence\"\n9\n10ln -sfn \"$DATA_PATH\" \"$HF_LEROBOT_HOME/physical-intelligence/libero\"
\n
\n
\n
\n \n

数据目录按 lerobot 格式组织：physical-intelligence/libero/{data,meta,videos}。norm stats 已随镜像提供于 assets/；若更换数据集需重算：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1python scripts/compute_norm_stats.py --config-name pi05_libero --num-workers 48
\n
\n
\n
\n \n

启动训练前，请确认 GPU 空闲、数据 / 权重已就位、磁盘可用空间充足，并为本次实验使用独立输出目录，避免覆盖历史实验。

\n

九、启动训练

\n

9.1 关键参数说明

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
参数值说明
--batch-size256（默认）全局 batch（跨 8 卡）
--num-workers8（优化）/ 2（默认）DataLoader worker 数，本负载主瓶颈
--fsdp-devices1（默认，纯 DP）>1 走 ZeRO-3 分片，省显存慢 6–8%；追吞吐峰值 / 多机弱扩展时使用
--model.max-token-len28（优化）/ 200（默认）libero 实测 max=21，降到 28 可省 FFN 计算
--weight-loader.params-path${BASE_CHECKPOINT_PATH}base 权重路径
--num-train-steps30000总步数
\n

9.2 优化开关

\n

以下开关默认关闭（等于上游行为），run_opt.sh 已默认开启全部优化：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1export OPENPI_SKIP_VIDEO_DECODE=1              # 跳过未消费的视频解码（占 71% 数据成本）\n2export OPENPI_SKIP_IMAGES=right_wrist_0_rgb    # 跳过恒 masked 的右腕图 token\n3export OPENPI_FLASH_ATTN=1 OPENPI_FLASH_IMPL=xla   # flash attention（XLA 实现）\n4export OPENPI_WORKER_CPU=1                      # 数据 worker 赶到 CPU，不抢 GPU\n5# CLI：--num-workers 8 --model.max-token-len 28\n6# config-B XLA flag（大 batch 下 +2.9%；OPENPI_STRICT_EQUIV=1 可关这组 flag）：\n7export JAX_OPTIMIZATION_LEVEL=O1\n8export XLA_FLAGS=\"--xla_gpu_triton_gemm_any=false --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_pipelined_all_gather=true\"
\n
\n
\n
\n \n

9.3 启动脚本

\n

容器默认工作目录 /root/workspace/openpi，三个脚本均在此目录：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1# 基线（上游默认：nw2 / tok200 / 无优化）\n2bash /root/workspace/openpi/run_baseline.sh\n3# 优化（全优化）\n4bash /root/workspace/openpi/run_opt.sh\n5# 追吞吐峰值（大 batch，需 fsdp 分片腾显存）\n6bash /root/workspace/openpi/run_opt.sh --fsdp-devices 8 --batch-size 2048
\n
\n
\n
\n \n

脚本支持 EXP_NAME=xxx、RESUME=1、SAVE_INTERVAL=n 等环境变量，以及任意 train.py 参数透传。

\n

9.4 断点续训

\n\n\n
\n
\n
\n Bash\n \n
\n
\n 
1# 首次训练（存 ckpt）\n2EXP_NAME=myrun SAVE_INTERVAL=1000 bash /root/workspace/openpi/run_opt.sh --num-train-steps 30000\n3# 中断后续训（同名 EXP_NAME + RESUME=1）\n4EXP_NAME=myrun RESUME=1 bash /root/workspace/openpi/run_opt.sh --num-train-steps 30000
\n
\n
\n
\n \n

十、多机训练（2 节点数据并行）

\n

openpi JAX 为 SPMD 单控制器，每节点启动 1 个进程执行 python scripts/train.py，该进程经 jax.jit + NamedSharding 驱动本节点全部 8 卡（不是每卡 1 进程）。各节点运行同一条命令，平台注入以下变量：

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
平台 env含义JAX 映射
MASTER_ADDR / MASTER_PORTcoordinator(rank0) 地址:端口coordinator_address
RANK节点 index（0=coordinator）process_id
WORLD_SIZE节点数（非卡数）num_processes
NPROC_PER_NODE每节点卡数仅信息用，JAX 自动探测本地卡
\n

启动脚本 run_multinode.sh 已封装上述变量读取、NCCL / IB 环境与优化配方。

\n

容器 RDMA / RoCE 要求（关键）：多机需容器能做 RDMA 内存注册，创建容器时必须具备以下能力：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1--network host --ipc host --gpus all \\\n2--device /dev/infiniband --cap-add=IPC_LOCK --ulimit memlock=-1 --ulimit stack=67108864
\n
\n
\n
\n \n\n

十一、性能参考

\n
\n

以下数据为特定测试环境下的参考值。实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。

\n
\n

单机 8 卡加速比（bs256 全局 batch，fsdp1）：

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
配置step time (s)samples/s
baseline20.1012.73
优化1.57163.01
加速比12.8×12.8×
\n
\n

baseline 吞吐偏低的根因是上游默认 num_workers=2 导致数据供给严重不足（GPU 几乎空等）；优化后转为 compute-bound。优化收益主力来自数据供给（worker 数提升 + worker CPU 隔离 + 跳过无用视频解码），计算侧优化（跳过 masked 图 token / flash-xla / 降低 max-token-len）为锦上添花。

\n
\n

单机吞吐峰值：

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
配置samples/s
优化 fsdp1 bs256（默认）163.01
优化 fsdp8 bs2048（追峰值）228.88
\n

双机 2×8 弱扩展：

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
配置samples/s扩展比
单机 fsdp8 bs2048228.881×（基准）
双机 fsdp16 bs4096443.411.94× / 96.9%
\n

等价性说明：优化版的 --model.max-token-len 28 与 OPENPI_SKIP_IMAGES=right_wrist_0_rgb 改变了输入张量形状，带来约 1e-4 的浮点漂移；两者数学上只影响被 mask 的无效 token，收敛终态一致，step0 两者仅差约 1e-4。数据顺序由 seed=42 固定、可复现。因此优化前后 loss 收敛曲线一致，训练结果等价。

\n

十二、常见问题

\n

Q1：多机训练 checkpoint 存盘卡在 Waiting for Save Finalize？

\n

checkpoint 目录必须落在所有节点同一路径可见的共享文件系统上。orbax 多进程 finalize 需跨进程协调 rename，若各节点写各自本地盘，finalize 屏障永不达成。请通过 --checkpoint-base-dir <共享存储路径> 指定共享存储目录。

\n

Q2：多机启动报 ibv_reg_mr_iova2 failed 或 NET/IB: No device found？

\n

容器缺少 RDMA 内存注册能力，请按「十、多机训练」中的要求加上 --device /dev/infiniband --cap-add=IPC_LOCK --ulimit memlock=-1。

\n

Q3：显存不足（OOM）？

\n

优先使用 --fsdp-devices >1 走 ZeRO-3 分片省显存（会慢 6–8%），或适当降低 --batch-size。

\n

Q4：离线环境下 lerobot 访问 HuggingFace Hub 报 Network unreachable？

\n

设置离线开关并确保数据软链已建立：

\n\n
\n
\n
\n Bash\n \n
\n
\n 
1export HF_HUB_OFFLINE=1\n2export TRANSFORMERS_OFFLINE=1
\n
\n
\n
\n \n

十三、后续操作

\n","fields":{"slug":"Dmvs45f6e","title":"在百舸上使用 openpi (JAX) 加速镜像训练 pi0.5 模型","date":"2026-10-28","extractedHeadings":[]},"headings":[{"value":"在百舸上使用 openpi (JAX) 加速镜像训练 pi0.5 模型","depth":1},{"value":"一、概述","depth":2},{"value":"二、适用场景","depth":2},{"value":"三、模型简介","depth":2},{"value":"四、环境与资源要求","depth":2},{"value":"五、准备工作","depth":2},{"value":"六、准备数据集","depth":2},{"value":"七、准备模型权重","depth":2},{"value":"八、配置训练环境变量","depth":2},{"value":"九、启动训练","depth":2},{"value":"9.1 关键参数说明","depth":3},{"value":"9.2 优化开关","depth":3},{"value":"9.3 启动脚本","depth":3},{"value":"9.4 断点续训","depth":3},{"value":"十、多机训练（2 节点数据并行）","depth":2},{"value":"十一、性能参考","depth":2},{"value":"十二、常见问题","depth":2},{"value":"十三、后续操作","depth":2}]}},"pageContext":{"isCreatedByStatefulCreatePages":false,"slug":"Dmvs45f6e","prev":{"id":"umvs23d15","name":"在百舸上使用lingbot-va加速镜像训练视觉-动作模型","path":"umvs23d15","filePath":"操作指南/快速开始/在百舸上使用lingbot-va加速镜像训练视觉-动作模型.md","seo":null,"parentIds":["ilib2qygp","3ma0icym2"],"parents":[{"id":"ilib2qygp","documentId":"bfa43a8b-968a-41a1-8c9d-906507eeaed9","name":"操作指南","repoName":"AIHC","filePath":"操作指南","disabled":false,"path":"ilib2qygp","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null},{"id":"3ma0icym2","documentId":"b5f72002-c562-4200-893d-0fdfd1ed1fe8","name":"快速开始","repoName":"AIHC","filePath":"操作指南/快速开始","disabled":false,"path":"3ma0icym2","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null}]},"next":{"id":"am794vnny","name":"全托管资源池","path":"am794vnny","filePath":"操作指南/全托管资源池/创建全托管资源池.md","seo":null,"parentIds":["ilib2qygp","Em794shv9"],"parents":[{"id":"ilib2qygp","documentId":"bfa43a8b-968a-41a1-8c9d-906507eeaed9","name":"操作指南","repoName":"AIHC","filePath":"操作指南","disabled":false,"path":"ilib2qygp","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null},{"id":"Em794shv9","documentId":"e6838c09-7abc-425b-bc8c-4800ee075222","name":"全托管资源池","repoName":"AIHC","filePath":"操作指南/全托管资源池","disabled":false,"path":"Em794shv9","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null}]},"parents":[{"id":"ilib2qygp","documentId":"bfa43a8b-968a-41a1-8c9d-906507eeaed9","name":"操作指南","repoName":"AIHC","filePath":"操作指南","disabled":false,"path":"ilib2qygp","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null},{"id":"3ma0icym2","documentId":"b5f72002-c562-4200-893d-0fdfd1ed1fe8","name":"快速开始","repoName":"AIHC","filePath":"操作指南/快速开始","disabled":false,"path":"3ma0icym2","lastMergeTime":null,"isApiDoc":null,"httpMethod":null,"seo":null,"sourceOrgName":null,"sourceRepoName":null,"sourceDocumentId":null,"solutionTag":null,"solutionColor":null}],"specificSeo":null}}}