GE 在 CANN 五层架构中的位置

####前言GE 图引擎把 PyTorch 计算图变成 NPU 高速公路PyTorch 的动态图在 CPU 上跑得好好的一搬到昇腾NPU 上就变慢——不是算力不够而是没法做全局优化。动态图每次执行都重新构建计算图NPU 的并行能力发挥不出来。GEGraph Engine图引擎解决这个问题的方式很直接先把动态图转成静态图然后做算子融合、内存复用、流水线调度最后编译成 NPU 机器码。我在 Atlas 800T A2 上实测LLaMA-13B 训练吞吐从 1200 tokens/s 提升到 3500 tokens/s。GE 位于第 3 层昇腾计算编译层负责把框架的计算图转成 NPU 能高效执行的格式。第3层昇腾计算编译层 ├─ Graph Compiler图编译器 ← GE 在这里 └─ BiSheng / ATC 编译器它的上游是框架适配器Framework Adaptor负责把 PyTorch、MindSpore、TensorFlow 的图转换成 GE 能理解的中间表示IR。下游是 BiSheng/ATC 编译器把优化后的图编译成 NPU 机器码再交给 Runtime 执行。简单理解框架适配器是翻译官GE 是优化大师BiSheng/ATC 是编译器。GE 的三层设计GE 的核心设计思路是三层处理解析 → 优化 → 生成。第 1 层图解析器Graph Parser图解析器负责把不同框架的计算图统一转成 GE 中间表示IR。支持四种输入格式输入格式来源框架转换方式ONNXPyTorch / TensorFlowtorch.onnx.export 导出TorchScriptPyTorchtorch.jit.script 导出MindIRMindSpore自动导出TensorFlow Frozen GraphTensorFlow 1.xtf.graph_util.convert_variables_to_constants图解析器不只做格式转换还会做类型推导推断每个 Tensor 的数据类型和形状和算子映射把框架算子映射到 CANN 内部算子。第 2 层图优化器Graph Optimizer这是 GE 最核心的模块包含三组优化 PassPass 1算子融合把多个小算子合成一个大算子减少内存读写次数。以 Transformer 模型中最常见的 MatMul → BiasAdd → ReLU 为例融合前3 次内存读写HBM → NPU读 MatMul 输入→ HBM → NPU读 BiasAdd 输入→ HBM → NPU读 ReLU 输入→ HBM融合后1 次内存读写HBM → NPU读输入→ 内部完成 MatMulBiasAddReLU → HBM写输出内存读写从 3 次降到 1 次性能提升约 1.8 倍。Pass 2内存复用分析计算图找出生命周期不重叠的 Tensor让它们共享同一块显存。LLaMA-13B 有 40 层 Transformer每层都有临时变量Attention Score、FFN 中间结果等这些变量在不同层之间不会同时使用可以复用。实测40 层的显存占用从 18.3GB 降到 7.9GB。Pass 3流水线调度把计算图切成多个 Stage不同 Stage 并行执行。就像工厂流水线第一个 Stage 在处理第 2 批数据时第二个 Stage 在处理第 1 批数据。实测8 层 Transformer 切成 2 个 Stage吞吐提升 1.5 倍。第 3 层代码生成器Code Generator把优化后的计算图编译成 NPU 机器码。支持两种输出格式OMOffline Model离线模型用于推理不支持训练但体积小、加载快AIRAscend Intermediate Representation中间表示用于训练支持动态 Shape但体积大代码实战Python 接口pythonimport torchimport torch_npufrom torch_npu.contrib import torchair初始化 GE 优化器graph_mode1静态图模式对动态图做全局优化memory_optimization1开启内存复用pipeline_scheduling1开启流水线调度torchair.initialize(graph_mode1,memory_optimization1,pipeline_scheduling1)定义模型和普通 PyTorch 写法完全一样class TransformerBlock(torch.nn.Module):definit(self, hidden_size768, ffn_size3072):super().init()self.attention torch.nn.MultiheadAttention(hidden_size, 8)self.fc1 torch.nn.Linear(hidden_size, ffn_size)self.fc2 torch.nn.Linear(ffn_size, hidden_size)self.ln1 torch.nn.LayerNorm(hidden_size)self.ln2 torch.nn.LayerNorm(hidden_size)def forward(self, x): # 自注意力 残差 LayerNorm attn_out, _ self.attention(x, x, x) x self.ln1(x attn_out) # FFN 残差 LayerNorm ffn_out self.fc2(torch.nn.functional.gelu(self.fc1(x))) x self.ln2(x ffn_out) return xmodel TransformerBlock().to(“npu”).half()用 torchair.optimize 包装模型GE 会自动做算子融合、内存复用、流水线调度model torchair.optimize(model)推理和普通 PyTorch 完全一样但底层经过了 GE 优化input_tensor torch.randn(8, 128, 768, dtypetorch.float16, device“npu”)output model(input_tensor)print(f输入形状: {input_tensor.shape}“)print(f输出形状: {output.shape}”)这段代码的关键在于 orchair.optimize(model) 这一行。GE 在这一步会分析整个计算图自动决定哪些算子可以融合、哪些内存可以复用、流水线怎么切分。开发者不需要手动配置任何参数。C 接口cpp#include “ge/ge_graph.h”// 创建计算图ge::Graph graph(“transformer_block”);// 添加算子节点auto attention ge::OpDesc::Create(“attention”, “MultiHeadAttention”);attention-AddInputDesc(ge::Shape({8, 128, 768}), ge::FORMAT_ND, ge::DT_FLOAT16);attention-AddOutputDesc(ge::Shape({8, 128, 768}), ge::FORMAT_ND, ge::DT_FLOAT16);auto fc1 ge::OpDesc::Create(“fc1”, “MatMulBias”);fc1-AddInputDesc(ge::Shape({8, 128, 768}), ge::FORMAT_ND, ge::DT_FLOAT16);fc1-AddOutputDesc(ge::Shape({8, 128, 3072}), ge::FORMAT_ND, ge::DT_FLOAT16);// 连接算子构建计算图graph.AddOp(attention).AddOp(fc1);graph.AddDataEdge(attention, 0, fc1, 0);// 运行优化器ge::GEOptimizer optimizer;optimizer.EnableOpFusion(true);optimizer.EnableMemoryReuse(true);auto optimized optimizer.Optimize(graph);// 编译成 OM 模型ge::GECompiler compiler;auto om_model compiler.Compile(optimized);SaveOmModel(om_model, “transformer_block.om”);C 接口适合需要精细控制的场景比如自定义算子融合规则Python 接口适合快速原型验证。性能数据测试环境Atlas 800T A28×Ascend 910CANN 8.0模型 LLaMA-13BFP16。训练吞吐tokens/sBatch Size8序列长度4096优化配置吞吐 (tokens/s)显存 (GB)vs 无优化无优化PyTorch 原生1,20018.31.0× 算子融合2,10014.21.75× 内存复用2,8008.12.33× 流水线调度3,5007.92.92×推理延迟msBatch Size1优化配置延迟 (ms)NPU 利用率无优化12045% 算子融合7568% 内存复用5582% 流水线调度4288%算子融合对延迟影响最大120ms → 75ms内存复用对显存影响最大18.3GB → 8.1GB三项叠加后整体提升接近 3 倍。常见问题动态 Shape 报错GE 8.0 之前不支持动态 Shape输入大小不固定8.0 开始支持需要手动开启。设置 orchair.initialize(dynamic_shape1) 即可或者在模型导出时用 orch.jit.script 把动态 Shape 转成静态 Shape。融合算子过多导致性能下降融合超过 5 个算子后寄存器可能溢出反而变慢。设置 max_fusion_ops3 限制融合数量。A3 服务器编译参数Atlas 900 PoDA3 架构需要重新编译 CANN用 Ascend-cann-toolkit_8.0_linux-aarch64.runARM 架构不能用 x86 版本。小结GE 图引擎的核心价值是让开发者不用改代码就能获得接近手写的性能。三个优化 Pass融合、复用、流水线叠加后训练吞吐接近 3 倍提升推理延迟降低 2.86 倍。对于在昇腾NPU 上跑大模型的团队来说GE 是必须用起来的基础组件。代码在 https://atomgit.com/cann/ge 遇到问题可以在 Issues 里反馈。