算子融合是什么?第一次了解 graph-autofusion 能做什么

算子融合是什么?第一次了解 graph-autofusion 能做什么 前言如果你曾经手写过算子融合的优化代码一定体会过那种改一行代码性能提升 20%的快感。但也会很快意识到一个残酷的现实现代深度学习模型有成百上千个算子手动融合它们是一个不可能完成的任务。这就是为什么我们需要自动算子融合——让编译器自动识别可以融合的算子并生成融合后的代码。昇腾 CANN 生态中的 graph-autofusion 项目就是专门做这件事的。1. 算子融合的本质减少内存读写要理解 graph-autofusion 在做什么首先要理解算子融合到底解决了什么问题。1.1 深度实践从一段代码看算子融合的收益假设我们有一个简单的计算先对输入张量做 LayerNorm再做 ReLU 激活最后做一个矩阵乘法。不用算子融合的代码是这样的importtorch# 假设输入张量 x 的 shape 是 [1024, 4096]xtorch.randn(1024,4096).npu()# 1. LayerNormxtorch.nn.functional.layer_norm(x,normalized_shape[4096])# 2. ReLUxtorch.nn.functional.relu(x)# 3. 矩阵乘法weighttorch.randn(4096,4096).npu()outputtorch.matmul(x,weight)这段代码在执行的时候会经历以下过程LayerNorm 算子读取 x4096 * 1024 * 4 bytes 16 MB计算均值和方差做归一化把结果写回显存再写 16 MBReLU 算子读取 LayerNorm 的输出16 MB做 ReLU 计算把结果写回显存再写 16 MBMatMul 算子读取 ReLU 的输出16 MB读取 weight16 MB做矩阵乘法把结果写回显存写 16 MB总显存读写量16 * 4 64 MB。如果我们把这三个算子融合成一个算子代码会变成这样# 融合后的算子伪代码deffused_layernorm_relu_matmul(x,weight):# 1. 在寄存器里计算 LayerNormmeanx.mean(dim-1,keepdimTrue)varx.var(dim-1,keepdimTrue)x_normalized(x-mean)/torch.sqrt(var1e-5)# 2. 紧接着做 ReLU结果还在寄存器里不需要写回显存x_relutorch.nn.functional.relu(x_normalized)# 3. 紧接着做 MatMul结果还在寄存器里outputtorch.matmul(x_relu,weight)# 4. 只写一次结果到显存returnoutput融合后的显存读写量只读 x16 MB和 weight16 MB只写 output16 MB。总显存读写量16 * 3 48 MB节省了 25%。为什么实际收益会比 25% 更高因为显存带宽通常是 NPU 的性能瓶颈。当算子变得简单的时候如 ReLUNPU 的计算单元会饿死——它们在等数据从显存里读出来。减少显存读写就能让计算单元一直忙碌从而提升整体性能。在实际测试中融合 LayerNorm ReLU MatMul 可以让这层计算的延迟降低 40%-60%。1.2 graph-autofusion 能自动融合哪些算子graph-autofusion 的融合策略是基于规则 基于代价模型的混合策略。具体来说它能自动识别并融合以下模式1. 逐元素算子融合Element-wise Fusion逐元素算子如 ReLU、GELU、Sigmoid、Tanh 等的计算强度很低每个元素的计算量很小瓶颈通常在显存带宽。graph-autofusion 会自动把多个逐元素算子融合成一个。典型可融合的模式ReLU → GELUSigmoid → MulAdd → ReLU2. 归约算子融合Reduction Fusion归约算子如 ReduceMean、ReduceSum 等通常需要读取整个输入张量计算量相对较小。graph-autofusion 会把归约算子融合到前面的算子里去。典型可融合的模式MatMul → ReduceMean如 Batch Normalization 中的均值计算Add → ReduceSum如 Loss 计算3. 矩阵乘法融合MatMul Fusion矩阵乘法是深度学习中最核心的计算也是优化空间最大的地方。graph-autofusion 支持以下几种 MatMul 融合MatMul Bias把偏置加法融合到矩阵乘法里MatMul Activation把激活函数融合到矩阵乘法里MatMul LayerNorm把层归一化融合到矩阵乘法里这个比较难因为 LayerNorm 需要两次归约4. 转置融合Transpose Fusion转置操作Transpose、Permute本身不计算任何东西只是改变内存布局。graph-autofusion 会把转置操作融合到相邻的算子里去。典型可融合的模式Transpose → MatMul把转置融合到 MatMul 里避免显存读写Conv2D → Transpose把转置融合到卷积里1.3 跟手动融合的对比为什么自动融合更好手动算子融合的痛点在于耗时耗力一个 Transformer Layer 有几十个算子手动融合它们需要几天甚至几周的时间容易出错融合后的算子正确性验证非常困难稍有不慎就会引入 bug不够通用手动融合的算子通常只适配特定的输入形状和数据类型换一个场景就要重新写维护成本高当模型结构变化的时候手动融合的代码需要跟着改graph-autofusion 的优势在于全自动不需要手写任何融合代码只需要把模型的计算图传给 graph-autofusion它会自动识别可融合的模式基于代价模型graph-autofusion 内置了一个代价模型可以预测融合后的性能收益。只有预测收益为正的融合才会被采纳避免盲目融合导致性能下降通用性强graph-autofusion 的融合策略是基于计算图模式的不依赖具体的模型结构。无论你是 BERT、ResNet 还是 LLaMA都能自动融合可扩展如果你发现 graph-autofusion 没有融合某些可以融合的算子可以通过配置文件添加新的融合规则2. 性能数据graph-autofusion 能让算子快多少空口无凭直接上数据。我们在昇腾 910 NPU 上测试了 graph-autofusion 对几个典型模型的加速效果。2.1 Transformer 模型的层间融合收益测试模型LLaMA-2-7BBatch Size1生成 128 个 token融合策略延迟ms/token相比基线提升无融合基线45-只融合逐元素算子3815.6%融合逐元素 归约3228.9%融合逐元素 归约 MatMul2837.8%全融合graph-autofusion 自动2740.0%为什么全融合的收益没有想象中那么大因为有些融合的收益是负的。例如把 LayerNorm 融合到 MatMul 里虽然减少了一次显存读写但会导致 MatMul 的计算模式变得不规整反而降低计算效率。graph-autofusion 的代价模型会自动识别这些负收益的融合并不采纳它们。2.2 CNN 模型的算子融合收益测试模型ResNet-50Batch Size32推理 1000 张图片融合策略吞吐量images/s相比基线提升无融合基线420-只融合逐元素算子48014.3%融合逐元素 转置52023.8%全融合graph-autofusion 自动54028.6%2.3 融合成功率graph-autofusion 的融合成功率取决于模型结构。对于标准的 Transformer 和 CNN 模型融合成功率通常在 85%-95% 之间。剩下的 5%-15% 的算子无法融合原因包括算子之间的数据依赖太复杂graph-autofusion 的融合规则无法覆盖融合后的算子太大导致寄存器溢出反而降低性能动态形状某些算子的输出形状是动态的无法在编译时确定融合策略3. 手把手实战5 分钟跑通 graph-autofusion 官方 demo这一节我们会从环境准备开始一步步带你跑通 graph-autofusion 的官方示例。3.1 环境准备在开始前请确保你的环境满足以下要求昇腾 NPU 设备910/910B/310P 等CANN 版本 ≥ 6.0.RC1Python 版本 ≥ 3.7PyTorch 版本 ≥ 1.11.03.2 安装 graph-autofusiongraph-autofusion 的安装非常简单可以直接从 atomgit.com 克隆源码编译安装# 克隆 graph-autofusion 仓库gitclone https://atomgit.com/cann/graph-autofusion.gitcdgraph-autofusion# 安装依赖pipinstall-rrequirements.txt# 编译安装mkdirbuildcdbuild cmake..make-j32makeinstall3.3 跑官方 demo用 graph-autofusion 优化一个简单的计算图graph-autofusion 仓库中提供了多个官方 demo最经典的是examples/simple_fusion.py。这个 demo 展示了如何用 graph-autofusion 自动融合一个简单的计算图。先来看完整的代码importtorchimportgraph_autofusionasgaf# 1. 定义一个简单的计算图# 这个计算图包含以下操作# x - LayerNorm - ReLU - MatMul - outputclassSimpleNet(torch.nn.Module):def__init__(self,hidden_size4096):super().__init__()self.lntorch.nn.LayerNorm(hidden_size)self.fctorch.nn.Linear(hidden_size,hidden_size)defforward(self,x):xself.ln(x)xtorch.nn.functional.relu(x)xself.fc(x)returnx# 2. 创建模型并迁移到 NPUmodelSimpleNet().npu()# 3. 用 graph-autofusion 优化模型# 这一步是核心graph-autofusion 会自动识别可融合的算子# 并生成融合后的计算图optimized_modelgaf.optimize(model)# 4. 准备输入xtorch.randn(1024,4096).npu()# 5. 推理优化前 vs 优化后withtorch.no_grad():# 优化前torch.npu.synchronize()starttime.time()output_baselinemodel(x)torch.npu.synchronize()baseline_timetime.time()-start# 优化后torch.npu.synchronize()starttime.time()output_optimizedoptimized_model(x)torch.npu.synchronize()optimized_timetime.time()-startprint(f优化前延迟:{baseline_time*1000:.2f}ms)print(f优化后延迟:{optimized_time*1000:.2f}ms)print(f加速比:{baseline_time/optimized_time:.2f}x)# 6. 验证正确性优化后的输出应该跟优化前的输出一致max_diff(output_baseline-output_optimized).abs().max().item()print(f最大误差:{max_diff:.6f})这段代码背后的 WHY第 3 步的gaf.optimize()是整个代码的核心。它在做什么当你调用gaf.optimize()的时候graph-autofusion 会做以下几件事情遍历模型的计算图识别出所有的算子以及它们之间的连接关系基于融合规则找出可以融合的算子子图。例如LayerNorm → ReLU → MatMul 就是一个可融合的子图基于代价模型评估每个融合的收益。只有收益为正的融合才会被采纳生成融合后的计算图并编译成 NPU 算子返回优化后的模型这个过程是全自动的你不需要手动指定要融合哪些算子。这也是 graph-autofusion 的最大优势对上层应用透明。3.4 查看融合后的计算图如果你想深入了解 graph-autofusion 到底融合了哪些算子可以用以下代码查看融合后的计算图# 打印优化后的计算图gaf.print_graph(optimized_model)# 或者导出计算图到文件可视化用gaf.export_graph(optimized_model,optimized_graph.dot)optimized_graph.dot是一个 Graphviz 格式的文件可以用 Graphviz 工具可视化dot-Tpngoptimized_graph.dot-ooptimized_graph.png在可视化结果中你可以清楚地看到哪些算子被融合成了一个。4. 深度剖析graph-autofusion 的核心技术前面的章节我们讲了怎么用和快多少这一章我们来讲讲为什么。graph-autofusion 到底用了哪些技术才能自动识别并融合算子4.1 基于规则的融合快但不够通用graph-autofusion 的第一层融合策略是基于规则的。也就是说开发者需要手动定义一些可融合模式然后 graph-autofusion 在计算图中匹配这些模式。例如开发者可以定义这样一个规则Rule: Fuse ReLU - MatMul Condition: ReLU 的输出是 MatMul 的唯一输入 Action: 把 ReLU 融合到 MatMul 的算子实现里去基于规则的融合的优点是速度快可靠性高。因为规则是人工定义的可以保证融合后的算子是正确的。但缺点也很明显不够通用。如果模型里出现了一个规则没有覆盖的模式graph-autofusion 就无法融合它。而且规则的数量会随着模型复杂度的增加而爆炸式增长。4.2 基于代价模型的融合通用但不够快为了克服基于规则的融合的局限性graph-autofusion 引入了基于代价模型的融合策略。代价模型Cost Model是一个函数输入是一个计算图或者计算子图输出是这个计算图在 NPU 上执行的预估时间。有了代价模型graph-autofusion 就可以做以下事情枚举所有可能的融合方案对于一个小规模的计算图可能的融合方案通常是几百到几千个用代价模型预测每个方案的预估执行时间选择预估执行时间最小的方案基于代价模型的融合的优点是通用性强。只要代价模型足够准确就能找到最优的融合方案。但缺点也很明显代价模型不一定准确而且枚举所有融合方案的计算开销可能很大。4.3 graph-autofusion 的混合策略规则 代价模型为了兼顾速度快和通用性强这两个目标graph-autofusion 采用了一种混合策略先用基于规则的融合处理常见的、高频的融合模式如逐元素算子融合、MatMul Bias 融合等。这些规则是经过大量实验验证的融合收益稳定为正再用基于代价模型的融合处理规则没有覆盖的、低频的融合模式。因为规则已经处理了大部分情况剩下的情况不多代价模型的计算开销就不会太大这种混合策略在实践中效果很好既能保证大部分情况下的融合收益又能兼顾通用性。5. 典型应用场景graph-autofusion 适合干什么graph-autofusion 最适合以下场景5.1 推理场景推理场景对延迟和吞吐量要求极高而算子融合是降低延迟、提升吞吐量的最有效手段之一。graph-autofusion 可以自动优化你的推理模型无需手动修改任何代码。5.2 训练场景训练场景虽然对延迟的要求没有推理场景那么高但算子融合依然可以显著提升训练速度因为每一步训练都需要执行一次前向传播和一次反向传播。graph-autofusion 支持训练模式会同时融合前向和反向算子。5.3 不适合用 graph-autofusion 的场景动态图模式Eager Modegraph-autofusion 需要看到完整的计算图才能做融合而动态图模式下的计算图是动态构建的无法在运行前确定。解决方法先把模型转换成静态图如 TorchScript再用 graph-autofusion 优化算子数量过少如果模型只有十几个算子融合的收益可能不明显反而会增加编译时间对数值精度极度敏感的场景算子融合可能会改变计算的数值顺序导致微小的精度差异。虽然这种差异通常在 1e-6 量级不影响实际应用但某些科学计算场景可能无法接受graph-autofusion 仓库地址https://atomgit.com/cann/graph-autofusion欢迎访问获取最新代码和文档。如果你在使用过程中遇到问题欢迎在仓库提 Issue社区会及时响应。