昇腾 A3 Pod 上 PI0模型训练适配案例

摘要

本文基于昇腾官方 Pi-0 训练文档与 PI0 训练记录，输出一份可复用的案例文档。案例覆盖环境与版本对齐、代码与依赖安装、数据集与权重准备（含受限网络/离线场景）、单机 8 卡训练验证（性能/精度）与日志指标口径；并补充DeepSpeed 多机训练、精度对齐方法（固定随机性、采集与比对精度数据）以及性能调优路径（profiling、融合优化器、绑核），用于支撑现场问题定位与验收复现。

业务背景以及面临挑战

业务背景

PI0（Pi-0 / pi-zero）是通用机器人基础模型之一，目标是训练出可在多机器人平台、多任务上泛化的通用机器人策略。期望在昇腾 NPU（A3 Pod）上完成训练跑通与验收，并在此基础上支持更长周期训练、精度对齐与性能调优，形成可复用方案。

硬件环境

软件版本

官方 Pi-0 文档给出的最低支持范围与版本对齐如下（以现场镜像为准）：

（DeepSpeed 训练分支）额外关注：

使用镜像

镜像：

pi0-pytorch2.1.0-ascend_cann8.2.rc1-py310-ubuntu20.04-arm64:0913

模型规格与参考

• 官方训练文档：Pi-0/README.md
• h20训练记录：pi0模型.md
• 权重参考（Pi-0）：https://huggingface.co/lerobot/pi0
• 上游仓库（参考实现）：https://github.com/huggingface/lerobot
• 昇腾适配参考（DrivingSDK）：https://gitcode.com/Ascend/DrivingSDK
• 文章参考：https://www.physicalintelligence.company/blog/pi0

PI0 模型说明

PI0（Pi-0 / pi-zero）是面向通用机器人控制的视觉-语言-动作（VLA）模型。其核心思路是利用预训练视觉语言模型（VLM）的语义与视觉理解能力，结合机器人数据学习到从“感知与指令”到“连续动作”的映射，从而在多任务、多机器人平台上具备较强的泛化能力。

模型关键特点：

• VLM 主干：基于预训练 VLM（PaliGemma），包含视觉编码器（SigLIP）与语言模型（Gemma），用于承载互联网规模预训练带来的语义与视觉知识迁移。
• 动作生成：使用 Flow Matching（扩散模型变体）将 VLM 输出扩展为连续动作生成，满足机器人控制需要的高频率连续输出。
• 动作专家模块：在 VLM 主干之外，引入更小的动作专家处理机器人状态与动作，形成感知-认知-行动闭环。
• 输入输出形式：输入通常包括最近多帧图像、文本指令与机器人本体信号（关节角等）；输出为未来一段时间的动作序列（action chunk），用于 50Hz 等实时控制场景。

面临的挑战/问题

• 受限网络/无外网环境下，训练代码可能无法自动下载依赖权重（例如 paligemma）。
• 依赖与版本不一致会导致运行失败或精度不一致（PyTorch/torch_npu/transformers/deepspeed 等）。
• 多机训练需要正确配置节点信息、主节点地址与 hostfile。
• 精度对齐需要可复现的随机性控制与可量化对比手段。
• 性能调优需要定位瓶颈（数据、计算、通信、调度）并验证收益。

问题解决方案与分析过程

解决方案和结果简述

问题根因

• 默认在线拉取权重与依赖，在无网或有防火墙环境下失败。
• 环境版本与 GPU 基线不一致导致行为差异或精度偏离。
• 多机分布式训练参数（MASTER_ADDR/PORT、NODE_RANK、hostfile、slots）配置不正确导致启动失败或资源利用不充分。

解决措施

• 训练前固化数据集与权重，必要时使用镜像站点下载；将本地 paligemma 路径写入代码，避免运行时在线下载。
• 统一版本对齐策略：优先对齐 transformers/deepspeed/torch 关键版本，并将改动收敛到脚本与可追溯配置。
• 多机训练使用环境变量注入 + 自动生成 hostfile 的方式，降低人工配置错误。
• 精度对齐按“粗对齐（代码/权重/数据/配置/超参）→精对齐（固定随机性、采集对比）”推进。
• 性能调优采用“采集 profiling → 分析瓶颈 → 单点优化验证收益 → 汇总复现”的闭环流程。

详细方案

方案 A：官方 Pi-0（Lerobot + Accelerate）训练跑通与验收

代码准备

按官方文档准备代码（commit 与补丁见 Pi-0/README.md）：

conda create -n pi0 python=3.10
conda activate pi0

git clone https://github.com/huggingface/lerobot.git
cd lerobot
git checkout 0cf864870cf29f4738d3ade893e6fd13fbd7cdb5

cp -f /path/to/PI0-TRAIN/Pi-0/pi0.patch .
cp -rf /path/to/PI0-TRAIN/Pi-0/test/ .
git apply pi0.patch

pip install -e '.[pi0]'
pip install -r /path/to/PI0-TRAIN/Pi-0/requirements.txt

数据集准备

• 数据集：koch_test（Hugging Face dataset）
• 下载完成后将数据集位置记作 {dataset_path}

受限网络建议使用镜像站点：

pip install -U huggingface_hub
export HF_ENDPOINT=https://hf-mirror.com

权重准备

1. 下载 Pi-0 预训练权重 {pi0_weights}

huggingface-cli download --resume-download --local-dir-use-symlinks False lerobot/pi0 --local-dir pi0

1. paligemma 权重 {paligemma_weights}

huggingface-cli download --resume-download --local-dir-use-symlinks False google/paligemma-3b-pt-224 --local-dir paligemma-3b-pt-224

下载完成后，将本地 paligemma 路径写入代码（脚本位于 paligemma_weights_mod.sh）：

bash test/paligemma_weights_mod.sh {paligemma_weights}

单机 8 卡训练验证

脚本位于 Pi-0/test：

• 性能（1000 steps）

bash test/train_8p_performance.sh {dataset_path} {pi0_weights}

• 精度（20000 steps）

bash test/train_8p.sh {dataset_path} {pi0_weights}

日志与指标口径（由脚本生成）：

• 训练日志：lerobot/test/output/train_*_8p_base_fp32.log
• 汇总指标：lerobot/test/output/Pi0_bs12_8p_acc.log
• 控制台输出：Final Performance images/sec、Final Train Loss、E2E Training Duration sec

官方参考结果

官方文档给出的单机 8 卡参考结果如下（Mixed 精度，global_batch_size=96，iterations=20k）：

实际运行的loss曲线：

方案 B：（DeepSpeed）多机训练、精度对齐与性能调优

本节内容来自记录（pi0模型.md），用于扩展到更长周期训练与多机场景。以下内容避免绑定任何企业内部路径，示例路径需按现场替换。

多机分布式训练要点（DeepSpeed）

关键参数：

• NNODES：节点总数
• NODE_RANK：当前节点编号
• MASTER_ADDR：主节点地址
• MASTER_PORT：主节点端口
• hostfile：各节点 IP 与 slots（每节点卡数/slot 数）

hostfile 示例：

192.168.1.1 slots=16
192.168.1.2 slots=16

常见做法是由平台注入环境变量（如 VC_WORKER_NUM/VC_TASK_INDEX）并从节点列表文件中读取主节点地址，自动生成 hostfile，减少人工错误。

精度对齐方法

对齐流程建议：

• 粗对齐：代码、权重、数据集、配置、超参、关键依赖版本（torch/deepspeed/transformers）
• 精对齐：固定随机性、采集精度数据、对比定位差异 API

固定随机性（NPU）可采用 msprobe 提供的 seed_all()（示例逻辑，实际插入位置以训练主入口为准）：

from msprobe.pytorch import seed_all
seed_all(seed=1234, mode=True, rm_dropout=True)

性能调优方法

性能分解：

单 batch 总时间 = 数据加载 + 模型前反向 + 优化器 + 后处理 + 通信 + 调度

建议闭环：

1. 采集 profiling（NPU profiler，导出到 tensorboard trace）
2. 使用 MindStudio Insight 导入分析，定位瓶颈与建议项
3. 单点优化验证（保持其余变量不变）

已验证的优化点与收益（以现场数据为准）：

• 算子/优化器：替换为昇腾亲和 AdamW（样本吞吐提升约 9.6%）
• 调度优化：开启粗粒度绑核（吞吐提升约 5.4%）

解决效果/价值

• 将官方 Pi-0 训练跑通路径与现场实战记录融合为一份标准化案例，可直接用于交付复现与验收。
• 覆盖离线权重准备、多机分布式训练、精度对齐与性能调优方法，降低现场问题定位成本。
• 输出统一的命令入口、日志位置与关键指标口径，便于回归对比与持续优化。

精度对齐解决效果

性能优化解决效果

优化前：

优化后：

参考

1. 昇腾 Pi-0 官方文档：Pi-0/README.md
2. DrivingSDK：https://gitcode.com/Ascend/DrivingSDK
3. Lerobot：https://github.com/huggingface/lerobot
4. PI0 Blog：https://www.physicalintelligence.company/blog/pi0