CANN π₀.₅模型训练优化说明

π₀.₅ 模型训练昇腾迁移与性能优化说明【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence背景介绍本案例在昇腾平台上使用 π₀.₅ 开源模型进行 LIBERO 基准的训练和评估并对训练过程的性能进行了深入分析和优化。以下内容将简要介绍 π₀.₅ 模型和 LIBERO 基准的背景信息并详细介绍在昇腾平台上训练时的性能分析与优化策略以及无 GPU 的仿真渲染方案。模型训练性能分析与优化在昇腾平台上训练pi05时profiling 最初暴露出四类比较典型的开销时间采样路径上的小算子与概率分布采样开销PaliGemma主干中 attention 前后的投影与算子调度碎片化训练图反向阶段的重复recomputation与 DDP 同步成本Gemma MLP 中的大matmul。优化项下表列出了当前代码中已经通过端到端训练验证的优化项。并已经在modeling_pi05.py、run_train.sh、run_profiling.sh和训练脚本默认配置中。优化项默认状态生效位置作用原理当前实测结论sample_beta默认开启sample_beta针对Beta(1.5, 1.0)默认时间采样避免通用分布采样实现带来的额外构造与调度开销保留对整体性能影响较小PI05_FUSE_PALIGEMMA_QKV1默认开启PaliGemmaattention 输入投影将q/k/v三个投影合并到一次线性层计算与拆分中减少 kernel 启动与中间张量调度保留attention 计算明显下降PI05_USE_NPU_FUSION_ATTENTION1默认开启PaliGemmaattention用torch_npu.npu_fusion_attention替换 eager attention减少 attention 主干中的算子碎片化保留训练性能提升约5%PI05_USE_NPU_GROUPED_GEMMA_INPROJ1默认开启Gemma MLPgate_proj/up_proj对共享同一输入的两段前向 matmul 使用npu_grouped_matmul合并执行反向保持标准线性层梯度公式保留训练级性能提升约1.0%~1.2%find_unused_parametersFalse、static_graphTrue、gradient_as_bucket_viewTruepi05默认开启DDP 包装层关闭无用参数探测固定训练图减少 bucket 重建和梯度通信额外开销保留但不是主要瓶颈--disable-outer-suffix-checkpoint推荐显式开启run_train.sh/run_profiling.sh关闭位置不合适的外层大 checkpoint避免反向阶段整段 suffix 被重复重算保留当前提升最大的单项性能优化第一类优化针对时间采样。pi05的默认时间分布是Beta(1.5, 1.0)原始实现走的是更通用的概率分布采样路径。对于训练中高频、参数固定的这一特例直接提供sample_beta可以减少不必要的分布构造和调度开销也避免在昇腾上退回到cpu进行计算。第二类优化针对 transformer 主干。PaliGemma侧的q/k/v投影默认改为融合执行同时训练阶段默认启用torch_npu.npu_fusion_attention。这样做的核心收益不是改变模型结构而是把原本多段 eager 投影、转置、attention 调度压缩成更适合昇腾执行的计算路径减少 kernel 启动次数和中间张量搬运。经过这一轮处理后attention 已经不再是训练的首要瓶颈。第三类优化来自训练图本身。对于pi05DDP 默认采用find_unused_parametersFalse、static_graphTrue和gradient_as_bucket_viewTrue目的是减少动态图额外探测、稳定 bucket 组织并降低梯度通信侧的额外开销。除此之外训练命令建议显式加上--disable-outer-suffix-checkpoint。原因在于 suffix 分支外层那一层大 checkpoint 会把整段后缀在反向时重新执行一遍而模型内部本来已经保留了更细粒度的 checkpoint两者叠加后反向阶段会出现位置不合适的重复重算。关闭这层外包式 checkpoint 后显存策略仍然可控但反向计算明显变快因此这是当前收益最直接的一项优化。第四类优化直接针对 profiling 中占比最高的 MLP 主热点。经过前几轮处理后训练热点已经集中到 Gemma MLP 的gate_proj/up_proj/down_proj三段矩阵乘法其中前两段具有完全相同的输入张量。当前版本在 NPU 上默认启用torch_npu.npu_grouped_matmul把gate_proj和up_proj合并成一次 grouped GEMM 前向计算由于该算子当前没有现成训练 autograd反向梯度采用显式手写的标准线性层梯度公式回传因此数值语义保持与原始GemmaMLP完全一致。这个优化之所以有效是因为pi05训练默认开启 gradient checkpointingMLP 前向在反向阶段还会被重算一遍因此只要把这两段前向 matmul 压缩成一次 grouped kernel就能在一个训练 step 内同时吃到“正向一次 反向重算一次”的收益。推荐训练与 profiling 方式下面给出当前保留优化对应的推荐训练与 profiling 命令。更完整的环境准备和脚本使用说明仍以上层 ../README.md 为准。训练推荐命令如下cd cann-recipes-embodied-intelligence/manipulation/pi05/train/src/scripts ./run_train.sh pi05 --nproc 2 --disable-outer-suffix-checkpoint如果需要重新采样 profiling推荐使用如下短窗口命令能够较快拿到稳定阶段的热点分布cd cann-recipes-embodied-intelligence/manipulation/pi05/train/src/scripts ./run_profiling.sh pi05 \ --nproc 2 \ --foreground \ --disable-outer-suffix-checkpoint \ --profile-wait 1 \ --profile-warmup 1 \ --profile-active 2 \ --profile-repeat 1从零复现推荐流程下面给出当前版本推荐的完整复现流程。该流程对应本仓库现有脚本默认设置。初始化代码与环境。cd your-workdir git clone https://gitcode.com/cann/cann-recipes-embodied-intelligence.git chmod x cann-recipes-embodied-intelligence/manipulation/pi05/train/src/scripts/setup.sh ./cann-recipes-embodied-intelligence/manipulation/pi05/train/src/scripts/setup.sh激活环境并进入脚本目录。conda activate lerobot cd your-workdir/cann-recipes-embodied-intelligence/manipulation/pi05/train/src/scripts首次训练前如模型权重和数据集已经缓存到本地可显式启用离线模式避免远端探测影响启动时间如果尚未缓存则不要打开这三个环境变量。export HF_HUB_OFFLINE1 export TRANSFORMERS_OFFLINE1 export HF_DATASETS_OFFLINE1启动推荐训练命令。./run_train.sh pi05 --nproc 2 --disable-outer-suffix-checkpoint如需快速采样 profiling使用下面的短窗口命令。./run_profiling.sh pi05 \ --nproc 2 \ --foreground \ --disable-outer-suffix-checkpoint \ --profile-wait 1 \ --profile-warmup 1 \ --profile-active 2 \ --profile-repeat 1训练与 profiling 结果检查。训练日志位于${PROJECT_ROOT}/ckpt/logs/train_model_timestamp.log。profiling 日志位于${PROJECT_ROOT}/ckpt/logs/profiling_model_timestamp.log。profiler 主目录位于本次训练output_dir下的profiling/。多卡训练产生profiling/rank0/...、profiling/rank1/...是正常现象。当前默认配置下pi05的稳定训练区间应大致落在5.75 ~ 5.85 s/it。Profiling 输出位置与结果解读profiling 日志位于${PROJECT_ROOT}/ckpt/logs/profiling_model_timestamp.logprofiler 主目录位于本次训练output_dir下的profiling/。多卡训练时看到profiling/rank0/...和profiling/rank1/...是正常现象表示每张卡各自产出一份 profiling 结果常用数据库文件位于profiling/rank*/.../ASCEND_PROFILER_OUTPUT/ascend_pytorch_profiler_*.db。仿真实时渲染由于 NPU 不支持 OpenGL 渲染因此在使用 NPU 进行仿真训练时需要将仿真环境的渲染模式切换为离屏渲染Offscreen Rendering。在使用 MuJoCo 仿真环境时可以通过设置环境变量来实现离屏渲染。具体设置如下Xvfb :1 -screen 0 1024x768x24 /tmp/xvfb.log 21 export DISPLAY:1 export LIBGL_ALWAYS_SOFTWARE1 export LD_PRELOAD/usr/lib/aarch64-linux-gnu/libOSMesa.so export LD_LIBRARY_PATH/lib/aarch64-linux-gnu/:$LD_LIBRARY_PATH export MUJOCO_GLosmesaMuJoCo CPU 渲染调用软件渲染OSMesa / EGL fallback在 CPU 上完成图像生成。Xvfb 虚拟显示创建一个虚拟显存帧缓冲模拟显示器环境让渲染画面有“输出窗口”。VNC 转发将虚拟显示画面通过 VNC Server 编码并转发用户可在本地 VNC 客户端实时查看。Citationmisc{cadene2024lerobot, author {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas}, title {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch}, howpublished \url{https://github.com/huggingface/lerobot}, year {2024} }【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考