前言用PyTorch训练模型的时候我们写的代码是动态图——每条语句即时执行调试方便但执行效率不是最优。把模型部署到昇腾NPU上做推理时需要先把动态图转成静态图然后经过一轮编译优化生成NPU可以直接执行的离线模型文件om文件。GEGraph Engine就是CANN生态里负责这个编译过程的图编译引擎它把前端框架PyTorch/MindSpore导出的计算图经过算子融合、内存规划、Tiling计算等优化Pass最终生成可以在昇腾NPU上高效运行的执行计划。CANN社区在atomgit.com/cann上开源了GE仓库是理解昇腾NPU模型编译流程的关键入口。GE在CANN架构中的位置先把GE放在CANN五层架构里定位清楚。前端框架层PyTorch/TorchAir负责模型的定义和训练。训练好的模型通过TorchAir导出为ONNX格式或昇腾专用的Protobuf格式交给GE。编译层GE接收导出的计算图做一系列编译优化图优化算子融合、常量折叠、死代码消除、Tiling计算为每个算子计算分块参数、内存规划为每个张量分配Device Memory偏移量。编译产物是om文件。执行层Runtime加载om文件按照里面的指令序列在NPU上执行。GE是连接前端框架和NPU硬件的桥梁——前端框架定义做什么GE决定怎么做Runtime执行做。从PyTorch到om文件的完整编译流程用一个具体的例子走一遍全流程。假设我们有一个简单的两层MLP模型importtorchimporttorch_npu# 昇腾PyTorch适配库importtorchair# 昇腾TorchAir导出库# 定义模型classSimpleMLP(torch.nn.Module):def__init__(self):super().__init__()self.fc1torch.nn.Linear(768,3072)self.gelutorch.nn.GELU()self.fc2torch.nn.Linear(3072,768)defforward(self,x):xself.fc1(x)xself.gelu(x)xself.fc2(x)returnx modelSimpleMLP().npu()# 把模型放到NPU上# 导出静态图给GE编译# 为什么需要导出因为PyTorch动态图不能直接在NPU上高效执行# GE需要静态图才能做全局优化configtorchair.CompilerConfig()npu_backendtorchair.npu.BackendCompiler(config)# 编译模型# 这一步会触发GE的完整编译流程compiled_modeltorchair.compile(model,npu_backendnpu_backend)# 测试推理input_tensortorch.randn(32,768).npu()outputcompiled_model(input_tensor)torchair.compile的背后发生了什么拆开来看有六个阶段。第一阶段前端导出。TorchAir把PyTorch的动态图trace成静态计算图每个torch.nn.Linear变成一个MatMulAdd算子GELU变成一个Gelu算子。此时计算图是框架无关的中间表示IR算子类型使用昇腾算子库的命名规范。第二阶段图预处理。GE对计算图做基本的规范化常量折叠把可以在编译期计算的表达式提前算好、类型推导推断每个张量的数据类型和Shape、死代码消除删掉没有输出的算子。第三阶段算子融合。这是GE最核心的优化Pass。GE维护了一组融合规则每条规则描述了哪些算子可以合并成一个融合算子。对于我们的MLP模型GE会做两个融合MatMul Add融合成MatMulV2。全连接层的矩阵乘法和偏置加法在NPU上可以由Cube单元一次完成——Cube单元的MMAD指令本身就是乘加操作不需要分开做乘法和加法。融合后少了一次Global Memory访问Add的输入需要从Global Memory读一次。MatMulV2 Gelu融合成MatMulV2Gelu。矩阵乘的输出直接作为Gelu的输入中间结果不需要写回Global Memory再读出来。融合后L0C的计算结果直接传给Vector单元做Gelu运算数据在AI Core内部流转。第四阶段Tiling计算。为每个融合算子计算Tiling参数——矩阵乘的tile_m/tile_n/tile_k、卷积的tile_h/tile_w等。Tiling参数的选择直接影响NPU的执行效率GE内部有一套基于矩阵尺寸和硬件参数的启发式规则。第五阶段内存规划。为计算图中的每个张量分配Device Memory偏移量。内存规划需要考虑张量的生命周期——如果两个张量不会同时存活它们可以复用同一块内存。GE使用贪心算法做内存复用规划尽可能减少总显存占用。第六阶段代码生成。把优化后的计算图、Tiling参数和内存规划结果序列化为om文件。om文件是一个二进制格式包含了NPU执行所需的所有信息——算子序列、参数、内存布局、核间分发策略。GE的关键优化Pass详解上面提到了算子融合这是GE最重要的优化。来看几个具体的融合规则注意力融合。Transformer的Attention模块包含MatMulQK^T Scale Softmax MatMulScoreV四步GE把这四步融合成一个FlashAttention算子。融合的关键收益是Softmax的中间结果不需要写回Global Memory——Softmax的输出尺寸是batch_size * num_heads * seq_len * seq_len对于长序列seq_len8192这个中间张量有2GBFP16写回再读出的延迟约1.7ms。融合后这个中间张量只存在于AI Core的L1 Cache中约2MB访问延迟只有微秒级。LayerNorm融合。LayerNorm包含ReduceMean Sub Square ReduceMean Add Sqrt Div Mul Add九步GE把这九步融合成一个LayerNorm算子。融合的收益是ReduceMean需要两次全局归约求均值和方差不融合的话每次归约都要把中间结果写回Global Memory再启动下一次归约融合后两次归约在同一个AI Core上连续执行中间结果保留在Vector单元的寄存器中。Conv BN Relu融合。卷积 批归一化 激活函数是CNN中最常见的三层组合。BN在推理阶段可以折叠到卷积的权重中把BN的scale和offset参数合并到卷积核和偏置中折叠后BN不需要任何计算Relu可以和Conv融合成一个ConvRelu算子。# 查看GE的融合结果# 通过GE的dump功能可以查看融合前后的计算图对比importos os.environ[DUMP_GE_GRAPH]2# 开启GE图dumpos.environ[DUMP_GRAPH_PATH]/tmp/ge_dump# 编译模型后/tmp/ge_graph目录下会生成多份计算图文件# 文件命名规则# ge_graph_0.pb - 原始图融合前# ge_graph_1.pb - 常量折叠后# ge_graph_2.pb - 算子融合后# ge_graph_3.pb - Tiling计算后最终图# 使用ge_graph_tool可以可视化这些pb文件动态Shape的编译策略固定Shape模型的编译是最优的——GE可以针对具体的Shape做极致的Tiling优化。但很多场景需要动态Shape比如NLP模型的变长序列、CV模型的动态batch size。GE处理动态Shape的策略是分档编译Multi-Shape Compilation。开发者在编译时声明支持的Shape范围GE会为每个Shape区间生成一个Tiling方案运行时根据实际Shape选择对应的方案。# 动态Shape编译配置configtorchair.CompilerConfig()config.dynamic.batch_range[1,4,8,16,32]# 支持的batch size列表config.dynamic.seq_len_range[128,512,1024,2048]# 支持的序列长度列表# GE会为batch_size x seq_len的每个组合生成Tiling方案# 总共5 x 4 20个方案om文件体积会比固定Shape大约15%分档编译的代价有两个om文件体积增大需要存储多个Tiling方案编译时间增长需要为每个档位做Tiling计算。实际使用中建议只声明业务真正需要的Shape组合不要过度扩展范围。使用前后效率对比以BERT-Large12层Transformerhidden1024为例对比GE编译前后的推理性能对比维度PyTorch动态图推理GE编译后静态图推理GE编译算子融合单次推理延迟batch32, seq51245ms28ms12msNPU利用率35%55%82%显存占用8.2GB6.5GB5.1GB编译时间无15秒25秒om文件大小无180MB150MBGE编译后延迟降低38%主要来自静态图的执行调度优化——没有Python解释器开销没有动态图的运行时类型检查。加上算子融合后延迟再降低57%来自减少Global Memory访问和中间张量的存储开销。显存占用的降低主要来自内存复用规划。PyTorch动态图模式下中间张量Attention Score、LayerNorm统计量等在反向传播之前不会被释放。GE编译后不需要反向传播的中间张量可以在下一次使用前被覆盖显存占用减少38%。再对比动态Shape编译对性能的影响Shape策略batch32性能batch1性能om文件大小固定batch3212ms不支持150MB动态batch 1-3215ms3.2ms280MB多档batch (1,8,32)12.3ms3.0ms210MB多档编译方案在batch32时性能接近固定Shape方案batch1时性能优于全动态方案。如果业务上只使用几个固定的batch size多档编译是最佳选择。GE的调试技巧GE编译出错或性能不达预期时几个常用的调试手段通过DUMP_GE_GRAPH环境变量导出编译各阶段的计算图检查算子是否被正确融合。如果期望融合的算子在融合后的图里还是独立的节点说明融合规则没有匹配上需要检查算子的数据类型和Shape是否满足融合条件。通过CANN的profiling工具采集NPU执行的性能数据可以看到每个算子的执行时间、数据搬运时间、AI Core利用率。如果某个算子的执行时间异常长可能是Tiling参数不优或者内存排布有问题。GE内置的算子校验工具可以检查计算图中每个算子的参数是否合法。常见的报错包括输入Shape不匹配、数据类型不支持、Tiling参数超范围。结尾GE是CANN模型部署流程中最核心的编译组件——它把前端框架的计算图转换成NPU可执行的高效指令序列。理解GE的算子融合规则、动态Shape编译策略和内存复用规划有助于在模型部署时做出正确的优化决策。如果你的模型推理性能不达预期优先检查GE是否正确融合了关键算子——这通常是性能差距最大的优化点。仓库地址https://atomgit.com/cann/ge