##
##前言
GE 图引擎：把 PyTorch 计算图变成 NPU 高速公路

PyTorch 的动态图在 CPU 上跑得好好的，一搬到昇腾NPU 上就变慢——不是算力不够，而是没法做全局优化。动态图每次执行都重新构建计算图，NPU 的并行能力发挥不出来。GE（Graph 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）。支持四种输入格式：

输入格式	来源框架	转换方式
ONNX	PyTorch / TensorFlow	torch.onnx.export 导出
TorchScript	PyTorch	torch.jit.script 导出
MindIR	MindSpore	自动导出
TensorFlow Frozen Graph	TensorFlow 1.x	tf.graph_util.convert_variables_to_constants

图解析器不只做格式转换，还会做类型推导（推断每个 Tensor 的数据类型和形状）和算子映射（把框架算子映射到 CANN 内部算子）。

第 2 层：图优化器（Graph Optimizer）

这是 GE 最核心的模块，包含三组优化 Pass：

Pass 1：算子融合

把多个小算子合成一个大算子，减少内存读写次数。以 Transformer 模型中最常见的 MatMul → BiasAdd → ReLU 为例：

`
融合前：3 次内存读写
HBM → NPU（读 MatMul 输入）→ HBM → NPU（读 BiasAdd 输入）→ HBM → NPU（读 ReLU 输入）→ HBM

融合后：1 次内存读写
HBM → NPU（读输入）→ 内部完成 MatMul+BiasAdd+ReLU → 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 机器码。支持两种输出格式：

• OM（Offline Model）：离线模型，用于推理（不支持训练，但体积小、加载快）
• AIR（Ascend Intermediate Representation）：中间表示，用于训练（支持动态 Shape，但体积大）

代码实战

Python 接口

`python
import torch
import torch_npu
from torch_npu.contrib import torchair

初始化 GE 优化器

graph_mode=1：静态图模式（对动态图做全局优化）

memory_optimization=1：开启内存复用

pipeline_scheduling=1：开启流水线调度

torchair.initialize(
graph_mode=1,
memory_optimization=1,
pipeline_scheduling=1
)

定义模型（和普通 PyTorch 写法完全一样）

class TransformerBlock(torch.nn.Module):
def init(self, hidden_size=768, ffn_size=3072):
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 x

model = TransformerBlock().to(“npu”).half()

用 torchair.optimize 包装模型

GE 会自动做算子融合、内存复用、流水线调度

model = torchair.optimize(model)

推理（和普通 PyTorch 完全一样，但底层经过了 GE 优化）

input_tensor = torch.randn(8, 128, 768, dtype=torch.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 A2（8×Ascend 910），CANN 8.0，模型 LLaMA-13B（FP16）。

训练吞吐（tokens/s，Batch Size=8，序列长度=4096）：

优化配置	吞吐 (tokens/s)	显存 (GB)	vs 无优化
无优化（PyTorch 原生）	1,200	18.3	1.0×
+ 算子融合	2,100	14.2	1.75×
+ 内存复用	2,800	8.1	2.33×
+ 流水线调度	3,500	7.9	2.92×

推理延迟（ms，Batch Size=1）：

优化配置	延迟 (ms)	NPU 利用率
无优化	120	45%
+ 算子融合	75	68%
+ 内存复用	55	82%
+ 流水线调度	42	88%

算子融合对延迟影响最大（120ms → 75ms），内存复用对显存影响最大（18.3GB → 8.1GB），三项叠加后整体提升接近 3 倍。

常见问题

动态 Shape 报错：GE 8.0 之前不支持动态 Shape（输入大小不固定），8.0 开始支持，需要手动开启。设置 orchair.initialize(dynamic_shape=1) 即可，或者在模型导出时用 orch.jit.script 把动态 Shape 转成静态 Shape。

融合算子过多导致性能下降：融合超过 5 个算子后，寄存器可能溢出，反而变慢。设置 max_fusion_ops=3 限制融合数量。

A3 服务器编译参数：Atlas 900 PoD（A3 架构）需要重新编译 CANN，用 Ascend-cann-toolkit_8.0_linux-aarch64.run（ARM 架构），不能用 x86 版本。

小结

GE 图引擎的核心价值是让开发者不用改代码就能获得接近手写的性能。三个优化 Pass（融合、复用、流水线）叠加后，训练吞吐接近 3 倍提升，推理延迟降低 2.86 倍。对于在昇腾NPU 上跑大模型的团队来说，GE 是必须用起来的基础组件。

代码在 https://atomgit.com/cann/ge ，遇到问题可以在 Issues 里反馈。