请添加图片描述

拿到一个 ONNX 文件,运行 atc --model=model.onnx --framework=5 --output=model,几十秒后拿到一个 .om 文件。这个过程里 ATC 到底做了什么?

先看输入。ONNX 是一个中间表示,存的是计算图的拓扑结构:节点是算子(Conv、ReLU、MatMul 等),边是 tensor 的流向。ATC 读进来之后,第一步是图解析——把 ONNX 的 proto 格式转成 GE 内部的图表示(ComputeGraph)。这一步会做算子映射:ONNX 的 Conv 对应 CANN 的 Convolution 算子,ONNX 的 Relu 对应 ReLU,大部分是一一对应的,少数需要拆开或合并。

映射完之后,第二步是图优化。GE 会跑一遍 pass 管线,包括:死代码消除(没有出度的节点删掉)、常量折叠(能静态算的常量节点直接算掉)、算子融合(Conv+BN+ReLU 合成一个 ConvBnRelu 算子,省掉中间结果的显存读写)。这套优化对性能影响很大——一个 ResNet50 的原始图有 300+ 个节点,优化后只剩 180 个左右,显存读写减少 40%。

优化完的图,第三步是内存规划。这是 ATC 最复杂的部分:给每个 tensor 分配显存偏移,目标是让生命周期不重叠的 tensor 复用同一块显存。比如 Conv 的输出 tensor,在下一层 ReLU 算完之后就不再用了,这块显存可以立刻分配给后面的 Pooling 层。GE 用的算法是「线性扫描 + 贪心合并」,显存复用率能做到 60-80%。如果模型很大(比如 LLaMA-70B),内存规划可能要跑几分钟。

第四步是算子选型。同一个算子(比如 MatMul),在昇腾上有多种实现:用 Cube Unit 做矩阵乘最快,但只支持特定 shape;用 Vector Unit 做更通用,但慢 30%。ATC 会枚举每个算子的所有实现,选一个最优的组合。这个过程叫「算子调度的静态绑定」,编译完就定死了,运行时不会再改。

第五步是代码生成。每个算子选中实现之后,ATC 生成对应的二进制指令(Cube 指令或 Vector 指令),打包进 .om 文件。这个文件里不光有指令,还有权重数据、内存规划表、算子执行顺序。加载 .om 的时候,Runtime 按表里的顺序把指令喂给 AICore,整个过程不需要再解释计算图。

实际编译时,有几个参数影响结果。--op_precision_mode 控制精度策略:设成 allow_fp32_to_fp16 的话,ATC 会把能转 FP16 的算子全转了,模型整体快 20-30%,但精度可能掉。--enable_fusion 控制算子融合,关掉的话编译快很多(10 秒搞定),但运行时慢 30-40%。--dynamic_batch_size 编译动态 batch,运行时 batch 在指定范围内不需要重编译,但 .om 文件会大 2-3 倍(因为每个 batch 都存了一份内存规划表)。

编译失败最常见的原因是算子不支持。ONNX 的算子有几百个,CANN 没必要全支持,只覆盖了 CV、NLP、推荐系统常用的那部分。遇到不支持的算子,有两个选择:第一,看这个算子能不能拆成已有算子的组合(比如 HardSigmoid 可以拆成 Clip + Mul);第二,用 Ascend C 写一个自定义算子,编译成 .so 之后注册到 OPP 目录,ATC 编译时就能找到了。

编译完的 .om 文件怎么用?aclmdLoadFromFile 加载,aclmdExecute 执行,aclmdUnload 卸载。用 PyTorch 的话不需要手动调这些 API,torch.npu 底层已经封装好了。但理解这个过程有助于排查问题:如果运行时报 `` 之类的错,大概率是 .om 文件和当前 CANN 版本不匹配,重新编译一下就好。

最后说一下 AOE 和 ATC 的关系。ATC 编译出来的 .om 里的算子参数是「通用最优」,AOE 做的是在这个基础上进一步搜索——比如 MatMul 的 tile 大小、卷积的循环展开次数,ATC 选的是保守值,AOE 会用 GA/RL 搜一遍,找到这个特定模型、特定输入 shape 下的最优值。搜完的结果存成一个调优记录文件,下次编译用 --load_tuning_result 加载,不需要重新搜。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐