使用 OpenPI-0.5 加速镜像进行训练
更新时间:2026-06-26
概述
本文介绍如何在百舸训练环境中,使用 openpi-0.5 加速镜像快速启动 pi05_libero 训练任务。该镜像已预置训练依赖,并内置了部分 OpenPI 基础资源,适合用于快速验证 OpenPI 0.5 在 LIBERO 数据集上的训练流程与性能表现。
功能介绍
- 预置训练环境:镜像内依赖已配置完成,启动任务后可直接使用镜像内 Python 环境运行训练。
- 内置基础资源:镜像内已缓存 OpenPI 相关资源,默认位置为
/root/workspace/openpi/.cache_openpi。 - 支持 LIBERO 训练:默认示例面向
pi05_libero任务与physical-intelligence/libero数据集。 - 训练性能优化:原始测试中,优化后 BS256 配置 step time 为 1.930s,samples/s 为 132.6,相比 baseline 提升约 78%。
计费说明
本文仅说明训练启动方式。实际费用以百舸训练任务运行时消耗的算力、存储等资源为准,请结合所在项目的资源与计费规则确认。
说明
开始前请确认具备以下权限或资源访问能力:
- 百舸项目中有可用 GPU 训练资源。
- 本快速开始说明为在双机八卡情况下进行训练的案例,您可以根据您的实际情况变更启动命令。
物料准备
数据集与权重已内置于镜像内,您可以直接使用,也可以参考使用格式替换为您自己的数据。
| 类型 | 说明 |
|---|---|
| 数据集 | huggingface 地址:https://huggingface.co/datasets/physical-intelligence/libero bos 下载地址:bos:/aihc-rdw-bj/huggingface.co/datasets/physical-intelligence/libero/ |
| 权重 | huggingface 地址:https://huggingface.co/lerobot/pi05_base bos 下贼地址:bos:/aihc-models-bj/lerobot/pi05_base |
| 镜像内置资源 | /root/workspace/openpi/.cache_openpi中:包含 gs://openpi-assets/checkpoints/pi05_base、gs://big_vision/paligemma_tokenizer.model、assets/pi05_libero/physical-intelligence/libero/norm_stats.json |
快速开始
1. 检查或生成归一化统计文件
镜像内已经对 pi05_libero 任务和 physical-intelligence/libero 数据集完成归一化,默认路径为:
Plain Text
1assets/pi05_libero/physical-intelligence/libero/norm_stats.json
如果需要重新生成,可执行:
Shell
1/root/workspace/openpi/.venv/bin/python3 scripts/compute_norm_stats.py --config-name pi05_libero
2. 启动训练
您可以按照实际情况修改以下参数:
模型权重
export PYTORCH_WEIGHT_PATH=
数据集
export HF_LEROBOT_HOME=
Shell
1#!/bin/bash
2
3export OPENPI_HOME=/root/workspace/openpi
4cd ${OPENPI_HOME}
5
6export HF_ENDPOINT=https://hf-mirror.com
7export UV_HTTP_TIMEOUT=60
8export HF_HUB_OFFLINE=1
9export TRANSFORMERS_OFFLINE=1
10export WANDB_MODE=offline
11export WANDB_DISABLED=true
12export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
13export OMP_NUM_THREADS=${OMP_NUM_THREADS:-8}
14export TORCHINDUCTOR_MIX_ORDER_REDUCTION=0
15
16export PYTHON="${OPENPI_HOME}/.venv/bin/python3"
17export OPENPI_DATA_HOME="${OPENPI_HOME}/.cache_openpi"
18export TORCHINDUCTOR_CACHE_DIR="${OPENPI_HOME}/torchinductor_cache"
19export ASSETS_BASE_DIR="${OPENPI_HOME}/assets"
20export CHECKPOINT_SAVE_BASEDIR="${OPENPI_HOME}/${TRAIN_RECIPE}_ckpt"
21
22export TRAIN_RECIPE="pi05_libero"
23export REPO_ID="physical-intelligence/libero"
24
25# 模型权重
26export PYTORCH_WEIGHT_PATH=
27
28# 数据集
29export HF_LEROBOT_HOME=
30
31default_norm_path="${ASSETS_BASE_DIR}/${TRAIN_RECIPE}/${REPO_ID}/norm_stats.json"
32if [ ! -f "$default_norm_path" ]; then
33 echo "norm_stats.json not found, computing..."
34 ${PYTHON} scripts/compute_norm_stats.py --config-name ${TRAIN_RECIPE}
35fi
36
37.venv/bin/torchrun --standalone --nnodes=1 --nproc_per_node=8 \
38 scripts/train_pytorch.py ${TRAIN_RECIPE} \
39 --exp-name ${TRAIN_RECIPE}_repro \
40 --batch-size 256 \
41 --num-train-steps 30 \
42 --log-interval 1 \
43 --save-interval 999999 \
44 --no-wandb-enabled \
45 --no-ddp-find-unused-parameters \
46 --ddp-bucket-cap-mb 25 \
47 --model.pytorch-compile-gemma-mode max-autotune-no-cudagraphs \
48 --pytorch-weight-path ${PYTORCH_WEIGHT_PATH} \
49 --assets-base-dir ${ASSETS_BASE_DIR} \
50 --checkpoint_base_dir ${CHECKPOINT_SAVE_BASEDIR}
性能参考
| 配置 | step time | samples/s | vs baseline |
|---|---|---|---|
| 默认 baseline | 3.442s | 74.4 | — |
| 优化后(BS256) | 1.930s | 132.6 | +78% |
| 优化后(BS384,max throughput) | 2.855s | 134.5 | +81% |
注意事项
- 不要默认使用
uv运行训练:镜像内依赖已配置完成,训练应优先使用镜像内.venv环境。 - 按实际资源调整并行度与 batch size:示例中
--nproc_per_node=8、--batch-size 256,请根据百舸任务实际 GPU 数量与显存调整。
误用 uv 后的环境恢复
如果不小心使用 uv 触发了依赖重装,可参考以下命令恢复环境:
Shell
1GIT_LFS_SKIP_SMUDGE=1 uv sync
2GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
3
4uv pip uninstall nvidia-nccl-cu12
5uv pip install --force-reinstall nvidia-nccl-cu12==2.27.5
6uv pip install cuda-bindings==12.9.4 cuda-pathfinder==1.2.2 cuda-toolkit==12.8.1 nvidia-nvshmem-cu12==3.4.5
7
8uv pip uninstall torch torchvision
9GIT_LFS_SKIP_SMUDGE=1 uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
10uv pip install -U triton==3.3.1
评价此篇文章
