背景
人形机器人和具身智能这两年热度很高,Pi0 这个视觉-语言-动作模型也经常被提到。做科研或者把模型往实际场景里推,算力是个绕不开的因素。最近刚好手头有台华为昇腾 Atlas 800I A2 服务器,我们试着把 Pi0 跑起来,顺便看看国产环境下的推理速度、控制精度和功能完整性到底怎么样。
CANN 是啥
CANN(Compute Architecture for Neural Networks)相当于昇腾芯片的软件栈,负责把上层框架的代码翻译成 NPU 能执行的算子。你可以把它想象成连接 PyTorch 和华为硬件的'翻译官'。没有它,代码根本没法在昇腾上跑。
这次环境搭建直接用了 CANN 开源社区的仓库 cann-recipes-embodied-intelligence,里面有针对具身智能场景的现成配置,省了不少自己对齐版本的时间。不过即便有仓库,环境配置的坑还是不少,后面细说。
环境搭建
(基于 cann-recipes-embodied-intelligence 仓库,按照 README 一步步来,但中间碰到几个依赖版本冲突。PyTorch 2.1 搭配 torch_npu 的版本必须严格匹配,否则一堆找不到符号的错误。另外 Pytorch 的 DDP 在 NPU 上需要调整通信后端,用 HCCL 替换 NCCL 就行。详细的步骤仓库里都有,这里不重复贴了,重点说说性能。)
推理性能
主要测了单次推理的耗时和误差。在 A2 上,Pi0 的一次前向推理大概 65 毫秒。这个数字是在 batch size 为 1、输入图像 256x256 的情况下测的,更复杂的场景还没试。对比一些云 GPU,差距不大,甚至某些算子比 V100 还快一点,可能是 CANN 对特定维度做了融合优化。
控制精度方面,端到端的动作输出误差在 1 厘米以内,对于抓取之类的任务基本够用。不过这个精度很依赖视觉输入的稳定程度,场景一暗或者物体纹理不明显,误差就会往上走。
功能与兼容性
Pi0 的模型结构没有变,所有标准组件都在 NPU 上跑通了。过程中发现一个细节:aten 算子中有一小部分需要在 CPU 上回退,但 CANN 社区已经把大部分都迁移好了。如果模型用了比较新的 Pytorch 语法,可能还是会碰到不支持的算子,但 Pi0 没这些问题。
总结
这次在昇腾 A2 上跑 Pi0 比预想的顺畅,国产硬件配合 CANN 软件栈已经能撑起具身智能的基本实验需求,65ms 的时延和 1cm 的误差也达到可以演示的水平。当然,大规模的 RL 训练稳定性和多机通信还没测,那是另一个话题。如果你也在用昇腾搭机器人实验,可以去 CANN 开源社区翻翻其他的优化案例,大模型和视觉方向的都不少。


