摘要

cann-recipes-embodied-intelligence仓库旨在针对具身智能业务中的典型模型、加速算法，提供基于CANN平台的优化样例，方便开发者简单、快速、高效地基于CANN平台使用具身智能模型。丰富CANN开源开放生态。

基于lerobot库，在昇腾Atlas A2环境适配pi0机器人端到端VLA具身智能大模型，通过使能昇腾融合算子、图模式、计算逻辑优化等手段，实现了较低的推理时延。单卡单次推理时间跑进80毫秒。目前，pi0模型在昇腾A2上的迁移优化代码已经开源在gitcode代码管理平台[https://gitcode.com/cann/cann-recipes-embodied-intelligence]。

欢迎大家多多star加星 / fork / 转发，多交流。


目录

  1. pi0机器人VLA大模型介绍
  2. pi0昇腾迁移优化目标
  3. pi0昇腾迁移适配方法
  4. 推理性能优化点分析
      4.1. pi0原始开箱推理profiling分析
      4.2. pi0推理代码结构昇腾潜在优化点分析
  5. 优化策略介绍
      5.1. 旋转编码sin/cos计算逻辑优化
      5.2. flash attention的npu原生算子优化
      5.3. layernorm部分npu原生算子优化
      5.4. qkv计算部分npu原生算子优化
      5.5. 图模式优化
  6. 优化策略结果总结
  7. 未来展望
  8. Citation

1. pi0机器人VLA大模型介绍

pi0是一个视觉-语言-动作（VLA）模型，专为通用机器人控制而设计。它基于预训练的视觉语言模型（VLM），结合流匹配（flow matching）机制，能够生成高频连续动作，实现对复杂、灵巧机器人任务的精准控制。整合OXE开源数据集和自有数据集，总计超过10,000小时机器人操作数据。在叠衣服、桌面清理、装盒等复杂任务上表现优异，零样本和微调设置下均显著优于现有基线方法（OpenVLA、Octo、ACT等）。成功完成5-20分钟的长时序多阶段任务，展现出强大的鲁棒性和泛化能力。pi0的模型架构图如下图 1-1所示：

pi0 模型架构
图 1-1　pi0 模型架构图

2. pi0昇腾迁移优化目标

在昇腾单机单卡计算平台上实现pi0模型的推理性能优化，保证推理精度达标，进一步缩短推理时间。测试的输入输出信息如下所示：

  • 输入：双路RGB图像 + 机器人本体状态 + 语言文本指令
  • 输出：50x6关节角度序列

3. pi0昇腾迁移适配方法

基于torch_npu，可以实现原始代码的昇腾NPU自动适配。同时，将代码中NPU不支持的float64数据类型替换为float32数据类型，可以实现推理快速适配打通，并保证推理精度达标。部分NPU迁移导入语句如下所示：

import torch_npu
from torch_npu.contrib import transfer_to_npu

4. 推理性能优化点分析

4.1. pi0原始开箱推理profiling分析

  • 为了分析pi0在昇腾A2服务器上原始开箱推理profiling性能，结合代码中torch_npu.profiler性能采集模块及MindStudio Insight计算流水性能可视化软件工具，得到pi0推理过程中的流水图及算子计算统计图。需要注意pi0在整个推理过程中，计算时间和空闲时间所占的比例。pi0原始开箱推理的流水图中，计算时间为1776.6毫秒，占比为22.5%，空闲时间占据的比例过大，为77.5%。由此可以得出基本是host-bound问题，即device侧一直在等待运算指令的下发。

  • 同时，也需要关注推理过程中AI_VECTOR_CORE和AI_CORE中算子所占的比例，pi0原始开箱推理过程中，AI_CORE算子占据的比例较小，为29.9%，由此可以得出计算过程中昇腾npu亲和的融合大算子所占比例较少，没能充分发挥NPU的Cube（矩阵单元）的计算性能，原始的pi0代码推理计算逻辑不亲合Cube计算单元。

  • 因此，需要着重减少pi0推理过程中CPU到NPU的指令下发次数，尽量使用昇腾融合大算子。同时，需要推动计算逻辑下沉至NPU侧，通过整图编译优化消除零散算子的调度开销，并针对旋转编码、Flash Attention、LayerNorm、QKV计算等关键性能瓶颈模块，逐一替换为昇腾亲和算子，最大化发挥AI_CORE的Cube矩阵计算能力。此外，还需优化CPU与NPU之间的数据传输链路，减少冗余数据交互，从“指令调度”“算子适配”“数据传输”等维度形成协同优化方案，解决host-bound问题与算力碎片化痛点，实现推理性能的质效提升。

4.2. pi0推理代码结构昇腾潜在优化点分析

根据pi0推理代码，慢系统调用PaliGemma的视觉编码器对单帧图像提取高维表征；快系统则是一个基于Transformer的Flow-Matching策略模型，负责并行去噪生成动作序列。原始lerobot仓中pytorch版本pi0代码逻辑主要存在以下两个潜在优化点：

  • 代码中的旋转编码部分、FA部分、layer-norm(add)部分、QKV计算部分均是通过小算子拼接或多步骤串接而成，没能利用昇腾融合算子的加速计算特性，导致在昇腾npu上计算不亲和。
  • 未启用图模式时，推理过程均是通过逐算子下发执行的方式，不仅会产生大量CPU与NPU之间的指令交互开销，还无法对零散算子进行融合编译优化，进一步加剧了host-bound问题（设备侧频繁等待CPU下发指令），同时小算子的串行执行模式也无法充分利用NPU的Cube矩阵计算单元，造成算力碎片化、空闲时间占比过高，严重制约整体推理效率。
  • 因此，从代码逻辑精简优化、昇腾亲和融合算子替换、图模式编译优化等手段出发，对pi0在昇腾NPU上的推理性能优化具有重要意义。以下是对几点优化策略的介绍。

5. 优化策略介绍

5.1. 旋转编码sin/cos计算逻辑优化

通过修改paligemma_with_expert.py中class PaliGemmaWithExpertModel(PreTrainedModel)中forward函数中的旋转编码计算逻辑，修改apply_rope函数中的输入输出，利用torch_npu.npu_rotary_mul融合算子，并将Q/K_states进行融合计算拆分，同时将sin/cos计算提到for循环外部进行统一计算，减少其重复计算次数等手段，可以实现旋转编码部分的优化加速。优化的部分代码片段如下所示：

def apply_rope(query_states, key_states, cos, sin):
  n_q = query_states.shape[2]
  n_k = key_states.shape[2]
  
  merged_states = torch.cat([query_states, key_states], dim=2)  # 维度为[B, S, n_q + n_k, D]
  
  merged_rot = torch_npu.npu_rotary_mul(merged_states, cos, sin)
  
  q_rot, k_rot = merged_rot.split([n_q, n_k], dim=2)

  return q_rot, k_rot

5.2. flash attention的npu原生算子优化

通过替换paligemma_with_expert.py中FA部分为npu_prompt_flash_attention融合算子，复用融合算子在npu上的快速计算能力，提升此模块的计算性能。同时，需要注意pi0中原始FA计算方式与npu_prompt_flash_attention融合算子在数学公式上的实现差别，尤其是在attention_mask上的数学公式实现差异（bool类型True位置转换为负无穷大），并进行相应的适配修改，将一些计算过程提到循环外面进行提前计算，可以实现FA部分的优化加速。优化的部分代码片段如下所示：

attention_mask = torch.logical_not(attention_mask).to(dtype=torch.int8, memory_format=torch.contiguous_format)
attention_mask = attention_mask[:, None, :, :]

......

def eager_attention_forward(
    self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
):
    att_output = torch_npu.npu_prompt_flash_attention(
        query_states,
        key_states.contiguous(),
        value_states.contiguous(),
        num_heads=self.num_att_heads_forward,
        input_layout="BSND",
        scale_value=self.scale_value,
        pre_tokens=65535,
        next_tokens=65535,
        atten_mask=attention_mask,
        num_key_value_heads=self.num_key_value_heads_forward
    )

    att_output = att_output.reshape(batch_size, -1, self.num_att_heads_forward * head_dim)

    return att_output

5.3. layernorm部分npu原生算子优化

将first residule计算中的layernorm和add运算替换为npu算子npu_add_rms_norm融合算子，复用融合算子在npu上的快速计算能力，提升此模块的计算性能。同时，需要注意pi0中layernorm计算方式和npu_add_rms_norm融合算子在数学公式乘积系数上的实现差别，并进行相应的适配修改。此外，可以将paligemma_with_expert.py中class PaliGemmaWithExpertModel(PreTrainedModel)中forward函数中原生的layer.input_layernorm计算部分用昇腾原生的npu_rms_norm进行替换，可以实现layernorm部分的优化加速。优化的部分代码片段如下所示：

hidden_states = torch_npu.npu_rms_norm(
    hidden_states,
    layer.input_layernorm.weight.add(self.ones_add),
    1e-6
)[0]

......

out_emb, _, after_first_residual = torch_npu.npu_add_rms_norm(
    layer.self_attn.o_proj(att_output[:, start:end]),
    hidden_states.to(torch.bfloat16),
    layer.post_attention_layernorm.weight.add(self.ones_add),
    1e-6
)

5.4. qkv计算部分npu原生算子优化

将pi0模型前向计算过程中q/k/v的分离计算过程，转换为qkv融合权重进行计算，然后再对融合计算结果进行split切分，节省缓存换入换出，用大矩阵乘法提升Cube的利用率。通过对paligemma_with_expert.py中的q/k/v计算部分进行对应优化，可以实现qkv计算部分的优化加速。优化的部分代码片段如下所示：

# 将q/k/v权重融合为单个线性层qkv
@torch.no_grad()
def fuse_qkv_weights(self):
    """
    将每层 self_attn 的 q/k/v 权重 concat 成单个 qkv 线性层
    """
    for model in self.models:  # paligemma + gemma_expert
        for layer in model.layers:      
            attn = layer.self_attn
            w_q = attn.q_proj.weight
            w_k = attn.k_proj.weight
            w_v = attn.v_proj.weight

            w_fused = torch.cat([w_q, w_k, w_v], dim=0).bfloat16()

            # 创建 qkv 线性层
            attn.qkv = nn.Linear(w_fused.shape[1], w_fused.shape[0], bias=False, device=w_q.device).to(torch.bfloat16)

            # 将拼接后的权重赋值给 qkv 权重              
            attn.qkv.weight.data.copy_(w_fused)

......

qkv = layer.self_attn.qkv(hidden_states)

q_proj, k_proj, v_proj = qkv.split([q_out, kv_out, kv_out], dim=-1)

5.5. 图模式优化

原生开箱的pi0推理性能，在npu上存在严重的host-bound问题。因此，为了减少CPU到NPU的逐算子下发与同步开销，结合torch_npu的torchair图模式编译模块，可以在modeling_pi0.py中将pi0的快系统与慢系统部分尽可能进行整图编译，进而实现整体的计算优化加速。优化的部分代码片段如下所示：

import torch_npu
import torchair as tng
from torchair.configs.compiler_config import CompilerConfig

config = CompilerConfig()
config.experimental_config.frozen_parameter = True
config.experimental_config.tiling_schedule_optimize = True

npu_backend = tng.get_npu_backend(compiler_config=config)

......

self.compiled_embed_and_model_forward = torch.compile(
    self.embed_and_model_forward,
    dynamic=False,
    fullgraph=True,
    backend=npu_backend
)

self.denoise_step_all_compile = torch.compile(
    self.denoise_step_all.forward,
    dynamic=False,
    fullgraph=True,
    backend=npu_backend
)

......

(
    prefix_pad_masks,
    dt,
    x_t,
    time,
    past_key_values
) = self.compiled_embed_and_model_forward(
    noise, images, img_masks, lang_tokens, lang_masks
)

for step in range(self.config.num_steps):
    self.denoise_step_all_compile(
        state,
        prefix_pad_masks,
        past_key_values,
        x_t,
        bsize,
        time,
        dt
    )

6. 优化策略结果总结

  • 基于上述昇腾环境上对pi0推理过程的协同优化，对单次推理过程进行基于MindStudio Insight的profiling分析。其中，昇腾npu计算时间总占比从22.5%提升至94%，计算总耗时从176.6毫秒降至71.6毫秒，NPU空闲时间从607毫秒降至4.58毫秒。AI_CORE类型昇腾融合算子计算时间占比从30%提升至50%。由此可以得出，优化后的pi0推理过程，已经极大地降低了host-bound问题，也更好地利用了昇腾融合算子加速能力。
  • 为进一步进行性能比较，避免profiling分析中的冷启动时延问题，进行100次连续推理计算时间取平均值的方式进行对比，得到昇腾A2优化后的pi0推理时间为80ms，相比于昇腾开箱670ms推理时延性能提升至8.3倍左右，提升至2.5倍左右的H20推理性能。
  • 分析结果显示出迁移优化策略提升了pi0对昇腾的亲和程度，发挥了昇腾npu的融合算子优势和算力优势，并充分显示出具身智能pi0模型在昇腾上迁移的可行性及易用性。为将来CANN开源开放生态中pi0等机器人具身大模型在昇腾系列处理器上的训练-推理-部署等提供了参考。

7. 未来展望

后续可以从以下几个方向，继续深化CANN开源开放生态中pi0等具身大模型在昇腾上的性能优化：

  • KV Cache 优化：针对pi0模型推理过程中 KV 缓存动态更新带来的性能损耗，采用固定尺寸缓存方案，按机器人控制任务的序列长度上限预设缓存空间，避免动态内存申请与数据拷贝，提升缓存利用效率。
  • 增量推理适配：基于固定KV Cache，构建pi0增量推理体系，针对机器人连续控制场景的增量输入（多帧图像、持续语言交互），复用历史计算结果，仅对新增特征片段做推理计算，进一步降低长时序任务的推理时延，适配实时控制需求。
  • 端到端全链路优化：将模型前后处理环节（图像解码、归一化、关节角度序列格式化）迁移至NPU执行，通过NPU原生图像处理算子替换CPU计算，消除数据传输延迟，进一步降低空闲时间，目标将推理总耗时压缩至更低。
  • 训练环节昇腾迁移：基于本次推理迁移经验，拓展至模型训练阶段的适配优化，利用昇腾NPU的分布式训练能力，实现Pi0模型的训练-推理全流程昇腾闭环，解决“训练在GPU、推理在NPU”的跨平台精度对齐问题。
  • 工具链与硬件兼容性提升：适配最新版本的昇腾工具链，优化代码适配逻辑，增强跨硬件型号的兼容性，同时推动优化策略的模块化封装，形成可一键调用的昇腾优化工具包，降低其他模型的迁移成本。
  • 最后，上述pi0适配昇腾平台并使能性能优化的相关代码，还有对应的昇腾环境搭建说明已在gitcode平台cann/cann-recipes-embodied-intelligence仓开源[https://gitcode.com/cann/cann-recipes-embodied-intelligence]，可以按照pi0昇腾迁移优化README文档进行使用。
  • 欢迎大家对本代码仓多多star加星 / fork / 转发，一起交流经验。

8. Citation

@misc{black2024pi0,
title={π0: A Vision-Language-Action Flow Model for General Robot Control},
author={Kevin Black et al.},
year={2024},
eprint={2410.24164},
archivePrefix={arXiv},
primaryClass={cs.RO},
url={https://arxiv.org/abs/2410.24164}
}

@misc{beyer2024paligemmaversatile3bvlm,
title={PaliGemma: A versatile 3B VLM for transfer},
author={Lucas Beyer et al.},
year={2024},
eprint={2407.07726},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2407.07726},
}

# python # 大模型精调推理 # 昇腾AI解决方案 # 人工智能深度学习 # 机器人
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

cover

昇腾NPU小模型推理性能调优实战：从1.5s到0.7s的优化之路

avatar 鲲鹏昇腾开发者社区
cover

【华为昇腾DVPP/AIPP学习篇】(3) AIPP+DVPP的使用

avatar 鲲鹏昇腾开发者社区

HDP 3.1 在鲲鹏服务器（CentOS 7.6 / openEuler 20.03）上的安装部署指南（优化版）

支持以下操作系统：CentOS 7.6bashbash说明：TaiShan 服务器支持的操作系统为 CentOS 7.6 / openEuler 20.03。

avatar 鲲鹏昇腾开发者社区