最近在 LoongForge 上完成了 Pi0.5 的sft加速训练 VLA 模型对我来说是一次陌生的尝试：3.5B 的模型、762 的序列长度，看上去用不了多少 trick。但是实际优化下来，发现可以细化的地方很多 最终 bf16下吞吐做到 24.0 samples/s（社区2.63 倍），FP8 下吞吐达到 27.34 samples/s（社区3.00倍） 1 冗余的256 序列长度：空相机 slot 裁剪 Pi0.5 的序列结构比较特殊，包含图像、文本指令和 action 三部分。为了理清 ViT 与主干 Attention 各自的计算负载，优化前先对各分支的序列长度做了下分解。 按两路相机的 LIBERO 数据推算，图像（2×256）、指令与 action（50）拼起来的联合序列本应是 762，但实际debug出来的序列长度却是 1018，平白多出了 256 的无效 padding。 排查发现 openpi 的 Pi0.5 输入固定 3 个相机 slot。当输入只有 2 路相机的时候，也会默认给缺的那路填一张占位图，再把 image mask 置为 False。 然而这个 mask 也只是让这 256 个 token 在注意力里不可见，实际的计算一点没少：占位图依旧需要计算SigLIP，产出的 256 个 token 也得过 18 层decode layer，中间的激活值照样保存，哪怕它们既不会被有效 token 注意到，也不参与最终的 loss 计算。 因此，LoongForge 在数据侧改为按需分配（将 data.num_images 设为实际相机数），只为真实存在的相输入生成slot，对于原本的pad token 不推进 position_ids ，裁剪后有效 token 的位置编码和可见范围不变，loss 与裁剪前严格一致，但是直接省去了约25%的计算量和相应的激活值显存占用。 2 密集的细碎算子开销：torch.compile 加速 通过 profile发现，Pi0.5 每层除了几颗大 GEMM，剩下全是小算子（尤其是 hidden size 较小的 action expert） ，包括 adaRMS 归一化、RoPE、两个分支 Q/K/V 的拼接与切回、gelu·up、门控残差 在eager 模式下，每个小算子都会触发一次独立的 kernel 启动，中间结果常驻显存。以 adaRMS 为例，cast 到 fp32、平方、求均值、rsqrt、chunk 出 scale/shift/gate、cast 回 bf16，每一步都产生一个完整大小的中间张量。 因此考虑进行torch.compile 加速，但是Pi05 这种比较动态的模型，不能对整个模型直接 compile， 一是vistion tower 是动态的，跟着图像数量走 embed_prefix 用 for 循环逐张图过 SigLIP，循环次数取决于数据集相机数 每张图的 pad mask 和块起始标记也是在 Python 里按图像数动态拼好，再通过 torch.tensor 搬到 GPU 上，直接编译极易触发断图（graph break）或重编。 二是 ViT 的耗时占比没有那么高，编译他它的收益相对 两路图像过 SigLIP 只占前向的 20%；18 层 decoder layer 占近 80% 因此，我们让Pi05 模型的 vision tower 和 prefix部分 拼接保持 eager ，只局部 编译后面形状固定的 backbone（ paligemma_with_expert ）开 torch.compile 之后，Inductor 主要做了三类融合： 逐元素 + 归约融成一颗 kernel ：adaRMS 整串操作融成一颗 persistent reduction kernel，均值和 rsqrt 单趟算完；残差加和 post_attention_layernorm 也融在一起 小算子融进 GEMM ：adaRMS 的调制和类型转换被吃进 q_proj 的 Triton GEMM，归一化结果不用单独落一次显存。 RoPE 整段折叠 ：cos/sin 表的构造（广播乘、cat、cos/sin、cast）融成一颗 kernel， q cos + rotate_half(q) sin 融成一颗，可以跑显存带宽上限 经过局部编译后， 训练吞吐提升了 54% 3 意外的显存收益: FlexAttention 压激活值 从LLM的惯性思维出发，我最开始看的是Pi05模型的Attention部分，但是profile发现，与普通的 LLM 不同，Pi0.5 的 attention 在端到端里的耗时占比并不高，主要问题是序列长度只有 762，attention 计算在 encoder layer 层内占比约 13%，剩下的主要是 MLP 占约 76%，折到端到端，attention 前向只占约 2.8% ，attention 部分的计算耗时优化空间很小 不过profile过程中发现，它的激活值多、算子碎。eager的实现每层先用 repeat_kv 把唯一的 KV 头复制成 8 份，再物化完整的 [B, 8, T, T] 概率矩阵；乘 scale、加 mask、两趟 softmax、PV、transpose 回去，每一步都是独立的 kernel，中间结果反向时都得留着，softmax 还会多存一份 fp32 中间结果，即使是torch.compile 包裹下，也没有显著降低激活值的开销 另一方面，Pi0.5 的掩码天然适合块稀疏（Block-Sparse）的 attention： Pi05的prefix（图像 + 文本）内部完全双向，但看不到 action，右上角整块屏蔽； action 能看到全部有效 prefix，50 个 action token 之间也互相可见； pad 所在的行和列整体屏蔽。 蓝色是 prefix 内部双向，浅橙是 action 看 prefix，深橙是 action 块内部，灰色是被屏蔽 掩码里大片区域要么全可见、要么全屏蔽，类似于 FlexAttention 的 BlockMask：按 128×128 的 tile 描述掩码，整块屏蔽的 tile 直接跳过，只有边界上的 tile 才逐元素判断。 因此替换 FlexAttn 后， enable_gqa=True 让 KV 头在 kernel 内广播，不再复制；反向只存输出和 logsumexp。速度和 eager 差不多（±2%），在E2E 更是区别不大，但 mbs32 下每层少存约 0.63 GiB，18 层合计约 11 GiB 4 显存和耗时的trade off: 选择性重计算 openpi 默认开启全量重计算：视觉塔、18 层 decoder layer、embedding 以及 action/time 投影的激活值一律不存，反向传播时全部重算一遍前向 在LoongForge框架下，叠加 torch.compile 和 FlexAttention 之后，虽然激活值显存占用已经大幅度下降，但是想要对齐官方的batch size，还是需要开启重计算 问题就变成了， 只重算哪一部分最划算 ？我们按模块逐个测算了重计算的“性价比”（即每省 1 GiB 激活需要多花多少毫秒）： ViT（SigLIP 27 层）：2.8 ms/GiB （+48.5 ms 换 17.5 GiB） ViT 的 hidden dim 较窄（1152 vs Gemma 的 2048），单位 FLOPs 产生的激活值最多，重算性价比最高。 PaliGemma MLP：4.0 ms/GiB （+150 ms 换 37.6 GiB）。 虽然激活占大头，但重算的绝对耗时太高。 Attention（norm + QKV + RoPE + FlexAttention）：7.8 ms/GiB （18 层合计仅 5.5 GiB）。 由于 MQA 的 K/V 很小，且 FlexAttention 不保存 score 矩阵，一旦重算却要把 QKV 投影和 attention 前向再跑一遍。 最终我们选择 只重算 ViT （通过LoongForge框架的 --activation-checkpoint-module-patterns 仅匹配 SigLIP 层），把重算开销降到了最低 5 无NVLink场景下的通信瓶颈: 分层编译隐藏梯度 AllReduce NVLink 场景下梯度通信通常不是瓶颈，但是在我们的实际训练的场景下， PCIe 通信代价高昂：Pi05 的3.5B 参数的 bf16 梯度约 7 GB，ring AllReduce 每卡需要穿书12.25 GB，通信耗时高达220ms。 然而profile 发现，由于一开始把整个 backbone 用torch.compile 编成一张图，Inductor 生成的反向是一个整体的 autograd Function，这导致所有参数梯度在反向结束时才同时 ready，DDP 梯度桶只能串行排在反向之后，导致每步有高达 ~215 ms 的通信开销完全暴露。 5.1 DDPOptimizer + torch.compile 预期用法：包住整个模型，和通信一起调度 这个问题 PyTorch 官方早就遇到过，其核心思路是： 编译器看到的范围越大越好 ——最好将整模型甚至整步训练交给它，由其统一调度梯度通信的时机与计算重叠。 DDPOptimizer （PyTorch 2.0 起，默认开启）：对 DDP 包住的模型调用 torch.compile 时，Dynamo 会顺着 DDP bucket 的大小将前向图切段分别编译，从而让反向梯度分批就绪。 但它的前提是 模型能被完整抓成整图 。 5.2 局部compile torch模型 和 DDPOptimizer 的不兼容 前面提到，Pi0.5 VLA 模型特殊的 ViT 结构决定了它无法对整个模型直接进行 torch.compile 编译只编 backbone 意味着在 DDP 内部对子模块显式 torch.compile ，这时 DDPOptimizer 会崩： 现象 ：清空编译缓存后第一次启动必挂，报 BackendCompilerFailed ，最终是 FakeTensor 上访问了 data_ptr() ；不清缓存直接重启又能跑。很有迷惑性，一度以为是 autotune 或 CUDA Graph 的问 根因 ： DDPOptimizer 在外面又套了一层，形成嵌套编译，DDPOptimizer 编译下一段子图前，会在 FakeTensorMode 下执行上一段已编译的子图来拿输出的 shape/stride，而里层 compile 的 AOTAutograd runtime wrapper 不知道自己处于 FakeTensorMode，会走真实的 autograd.Function.apply ，碰到 FakeTensor 的 data_ptr() 就挂。 热缓存能跑，只是因为命中 FX graph cache 绕开了这条路径，换个 shape、换个 rank 数或者清一次缓存就会复现。 所以只能先关掉它（ --no-dynamo-optimize-ddp ）。 代价正如前面所分析的：由于无人帮忙切图，整张反向图跑完后梯度才能 ready ，通信完全暴露出去了 5.3 自行分层编译，手动切图 既然整图编译无法实现、 DDPOptimizer 又无法使用，干脆就在 backbone 内部 自己动手切图 ： 我们将编译策略改为 multi_group 分组编译——18 个 decoder layer 各自作为一个独立的编译单元（相同 shape 只需编译一次，其余 17 层直接复用），同时将 final norm、action 输入投影 + time MLP、action 输出投影各自独立成小图。 这样一来，每一层的反向都成为了独立的 autograd.Function ：当第 N 层的反向一结束，它的梯度桶就能立刻发起 AllReduce，与第 N-1 层的反向计算实现并行重叠。 本质上，这相当于 手动实现了 DDPOptimizer 想做的事 ，只是将切分点精准锁在了层边界上，完全摆脱了对 DDPOptimizer 的依赖。改用手动分层编译后，暴露的通信基本消失，吞吐+10% 6 bf16 效果验证：LIBERO 测评 经上述一系列 bf16 优化叠加，整体训练吞吐成功飙升至 24.0 samples/s ，达到社区吞吐的 2.63 倍 。 效果表现方面，LoongForge bf16 模型在 LIBERO 四个子集上的 5 轮平均测评结果如下： 子集 社区 复现社区 LoongForge （bf16） libero_spatial 97.0 96.0（−1.0） 99.6（+2.6） libero_object 99.0 98.0（−1.0） 99.6（+0.6） libero_goal 98.0 99.0（+1.0） 98.8（+0.8） libero_10 (Long) 96.0 96.0（0.0） 97.0（+1.0） 平均 97.50 97.25（−0.25） 98.75（+1.25） 括号里是相对社区的差值，LoongForge 版本平均 98.75，比社区高 1.25 个点，四个子集均不低于社区。训练加速下，模型测评效果符合预期 7 极致的 FP8 计算优化 到这里， bf16 下能抠的性能基本已经见底。接下来，我们将 Pi0.5 中的 Linear 模块全面推进到 FP8 计算：输入与权重量化为 FP8 后执行 GEMM，结果按 scale 还原，计算前后保持 bf16 7.1 TransformerEngine 的 FP8 尝试 最先尝试的是 TransformerEngine的FP8 方案：把 nn.Linear 换成 te.Linear ，前向套上 fp8_autocast ，但是真正的问题是它和前面的优化冲突： 和 torch.compile 冲突，到处断图 ： TE 内部有被 @torch.compiler.disable() 装饰的函数（fp8 元数据和 amax 处理） Dynamo 于是在每个 te.Linear 处断图，218 个

  • 情报分类：服务器与云资源
  • 分类依据：内容涉及服务器、云资源或网络线路
  • 信息来源：服务器 / LINUX DO - 最新话题
  • 发布时间：2026/10/19 22:34:53