☰
深入CANN metadef:内存复用协议如何压榨显存峰值
2026/10/20 4:49:55 网站建设 项目流程

开场:一个困扰分布式训练团队很久的问题

做AI框架底层开发的朋友,对“显存不够用”这件事应该都不陌生。无论你是跑大模型微调,还是做多路视频推理,经常遇到的情况是:明明单算子显存占用都不大,模型整体跑起来却把几十G的显存吃得干干净净。更气人的是,同样一个网络结构,换一套编译选项,内存占用能差出20%到30%。这背后的差距,往往不在于算子本身实现得好不好,而在于框架调度层有没有做一件事——内存复用。

我在深入阅读CANN(Compute Architecture for Neural Networks)的源码和架构文档时,发现它的内存管理设计里藏着一个非常关键的层次,叫metadef,全称是“内存元数据定义层”。这层负责把所有算子执行时需要用到的内存信息抽象成元数据,再基于这些元数据做图级别的内存复用规划。可以说,你看到的显存占用数字,很大一部分不是由算子决定的,而是由metadef这一层怎么构建内存图、怎么设计复用协议决定的。

这篇文章,我想把这层的东西拆开讲清楚。适合正在做AI编译器、算子适配、推理引擎优化,或者单纯想搞清楚自家模型显存为什么爆掉的同学看。不需要你之前接触过CANN,只要你写过CUDA或者PyTorch自定义算子,理解起来都没问题。

1. 为什么需要一层“内存元数据定义”

1.1 从GE的算子下发逻辑说起

在CANN的架构里,用户侧写好的网络先经过图引擎(Graph Engine,简称GE)做解析、优化、融合,之后会生成一个又一个的Task,这些Task最终被下发到昇腾芯片上执行。但这里有个核心问题:每个Task执行时都需要一块内存来放输入、输出和中间结果。如果每个Task都按最大需求向系统申请独立内存,那模型占用的显存就是所有算子峰值之和——这在任何真实场景下都是不可接受的。

早期版本的GE确实走过这种“实打实申请”的老路。算子多、模型深的时候,显存占用线性上涨,一旦网络里出现几个特别“胖”的算子(比如大矩阵乘法、大卷积),整个显存直接被打满。后来GE架构演进时,设计者意识到:必须在“执行层”之上抽象出一个“内存描述层”,先把“谁需要内存、需要多大、什么时候开始需要、什么时候可以释放”这件事从业务逻辑里抽离出来,独立建模。这个抽象层,就是metadef。

把内存管理从执行流程里抽出来,换来的是两样东西:一是规划能力,我们能提前为整张图做一次性内存规划,而不是执行到哪算到哪;二是复用能力,内存块的分配和回收可以基于图的拓扑信息做全局优化。这两样东西,正是内存复用协议能够成立的基础。

1.2 metadef到底包含哪些内容

metadef这个名字听起来抽象,其内部实际承载的数据结构非常具体。它主要由三个子定义域构成:图定义(GraphDef)、内存定义(MemDef)和模型定义(ModelDef)。GraphDef把GE优化后的计算图转成一种更底层的、面向内存规划的格式;MemDef则定义了内存块、内存区、内存偏移这些核心概念;ModelDef负责把整张计算图的内存规划结果封装成一个可以加载执行的模型对象。

打个比方:GE像是一个城市的规划局,决定路网怎么布局、交通流怎么走;metadef则像是国土详测部门,把每一块地皮的面积、边界、用途登记成册。算子真正跑起来的时候,其实是在“使用”由metadef规划好的地块,本身不需要关心地块是怎么划分的。

从这个意义上说,metadef的核心价值是“解耦”。它让算子开发者和框架调度者之间形成了一道清晰的分工线:算子开发者只需要描述“我这个op需要几块内存、每块多大、对齐要求是什么”,框架调度者则根据整个图的拓扑来决定“这些内存块能不能复用、怎么复用”。

2. 图构建链路:从GE优化图到内存元数据图

2.1 计算图编译的三个典型阶段

实践中,一张用户写好的神经网络模型,在CANN里要经过三个典型阶段才会变成可执行的任务。

第一个阶段是GE的构图阶段。这个阶段把用户侧框架(比如TensorFlow、PyTorch)传来的计算图做基础优化:算子消除、常量折叠、数据格式转换、算子融合。融合操作非常关键,比如把Conv、BN、ReLU三个算子融合成一个算子,不仅减少了kernel launch次数,也直接减少了内存块数量。内存复用的基础是内存块的生命周期清晰可控,算子融合恰恰让很多原本独立的中间结果直接被消灭在寄存器或片上缓冲区里。

第二个阶段是metadef的GraphDef构建阶段。GE优化后的图会通过一个转换器,生成metadef内部的GraphDef结构。这里不是简单地把图重新画一遍,而是对每个节点做“内存特征提取”——记录这个节点的输入张量形状、输出张量形状、数据格式、对齐要求,以及它内部可能需要的工作空间(workspace)大小。例如一个Conv2D算子,除了输入输出外,在特定算法下可能还需要一块临时空间做im2col或winograd变换,这块空间的大小也需要在这个阶段被登记。

第三个阶段才是内存规划。拿到GraphDef之后,metadef会对整张图的所有内存需求做全局分析,决定哪些内存块可以复用同一片物理内存,最终输出MemDef。

这里最容易踩的坑是:有人直接在GE图上做内存规划,跳过GraphDef这一层。短期看能省掉一层转换,但长远看问题很大。一个算子在GE图里可能是一个节点,到了实际执行时可能被拆成多个子任务,每个子任务的内存需求不同;反过来,多个GE节点可能融合成一个执行算子。没有GraphDef做一次统一的“内存语义归一”,内存规划就会建立在错误的基础设施上。

2.2 节点、边与内存节点的生成规则

GraphDef的构建有一个重要特点:它的节点并不完全等价于GE图中的算子节点。GraphDef除了保留算子节点本身外,还会额外生成两类“内存相关节点”:一类是张量节点(Tensor Node),用来描述一个具体的中间张量;另一类是内存申请节点(Allocate Node),用来描述在某一次执行中需要申请的一块内存区域。

之所以要区分张量节点和内存申请节点,是因为Tensor的生命周期和Memory的生命周期并不是一一对应的。同一个张量在不同执行阶段可能被多个算子消费,而每份内存只能被物理占用一次。把概念分开之后,规划逻辑就清晰了——先做Tensor级别的生命周期分析,再做Memory级别的复用匹配。

构建过程其实是从输出节点反向遍历的。从网络的loss节点开始,沿着数据依赖边反向遍历每个算子,为每个算子的输出张量建立一个Tensor Node,并把这个张量的shape、dtype、format、alignment信息挂到节点上。同时为每个算子建立一个AllocateNode,登记它的workspace需求。

值得强调的是反向遍历这个细节。反向遍历是为了保证“使用顺序”信息完整:一个张量如果被多个下游算子使用,那么它的存活期必须覆盖到最后一个使用它的算子执行完。用反向遍历能自然地聚合这种依赖信息。如果正向遍历,你需要在后续步骤里额外做数据流分析才能补齐这些信息。

2.3 内存元数据的数据结构示例

实际代码里,GraphDef的内存节点通常长这样(下面代码是示意写法,用于说明核心字段):

struct TensorDesc { std::string name; std::vector<int64_t> shape; DataType dtype; Format format; uint64_t align; // 对齐字节数 }; struct OpMemDesc { std::string op_name; std::vector<int64_t> input_ids; std::vector<int64_t> output_ids; int64_t workspace_size; // 算子内部需要的临时空间 uint64_t workspace_align; }; struct AllocateNode { int64_t tensor_id_or_ws; uint64_t size; uint64_t alignment; uint64_t lifetime_start; // 生命周期起点,执行序号 uint64_t lifetime_end; // 生命周期终点 };

构建完成后,每个内存持有者都有一个清晰的lifetime区间。这个区间就是后续内存复用决策最核心的输入。

3. 内存复用协议:复用的时机与规则

3.1 生命周期分析:内存复用的基石

内存复用协议要回答的核心问题只有一个:两块内存能不能放在同一片物理地址上?

答案是:只要两块内存的存活期不重叠,它们就可以共享同一片物理内存区域。这就是“时间换空间”的核心逻辑。假如算子A的输出张量在算子C执行完后才被最终消费,那么这块内存从A结束到C开始的这段时间段内就是“空闲”的。这期间如果有另一个算子B需要一块同样大小的内存,完全可以把B的内存地址映射到A的内存区域上。

在这个框架里,生命周期分析的具体做法是先给所有结点做一次拓扑排序,然后为每个结点分配一个“执行序号”。执行序号是单调递增的整数。某块内存的lifetime_start就是创建它的算子执行完成后的序号,lifetime_end则是最后一个读取它的算子执行开始前的序号。

在真实源码中,这个分析过程还带有一些细节。比如一个算子的输入张量可能是另一个算子的输出张量,那么这个输入张量的lifetime_start必须晚于产出它的算子的执行完成时刻。又比如,同一个算子的多个输出张量,它们虽然产出时刻相同,但由于被不同下游消费,lifetime_end各不相同。这意味着它们能参与的复用集合也不一样,规划时需要区分对待。

3.2 离散内存池与连续内存池的选择

做内存复用,第一个要决策的问题是:复用的物理单位到底是一整块显存,还是细粒度的离散小块?

从metadef的设计来看,它同时支持两种策略。第一种是离散内存池(Discrete Memory Pool),适合网络中存在大量大小不一、生命周期散乱的中间张量的情况。系统维护一个空闲块列表,每次有内存需求时,从空闲块里找一个大小满足要求、地址对齐满足要求的块分配出去;释放时把块重新挂回空闲列表。这种方式灵活,但存在碎片问题——这也解释了为什么很多模型跑久了之后显存占用越来越高。

第二种是连续内存池(Continuous Memory Pool),也叫静态规划。它在图编译阶段一次性算出总需求峰值,然后分配一整块连续内存。每个张量在这块连续区域里占据一个偏移量,偏移量一旦确定就不再变化。这种方式几乎没有运行时的分配开销,也不用担心运行时的碎片问题,是静态图推理引擎最常用的方案。或者更准确地说,是离线编译场景下最稳妥的方案。

那么,metadef怎么在这两种策略之间做取舍?实际设计是按场景来分的。对于在线推理场景,执行图相对固定,内存需求可以离线预计算,使用连续内存池能最大化确定性;对于训练场景,由于存在动态shape、控制流、算子重编译,内存需求往往是运行时才知道的,这时候离散内存池更合适。

3.3 复用匹配算法:贪心还是暴力

确定了“可以复用”的标准之后,真正的实现问题来了:给定N个张量的生命周期区间和大小,怎么把整体内存峰值压到最低?

如果张量数量少,当然可以用二分图匹配甚至回溯法找全局最优解。但真实模型里的中间张量动辄成千上万,暴力搜索直接爆炸。因此metadef默认采用的是一套贪心式复用策略:把所有内存需求按lifetime_start排序,依次分配;每次分配时,在“当前已经空闲的内存块”里找一个大小大于等于需求、且地址对齐满足要求的最小块(best-fit)分配给新的张量。

为什么是best-fit而不是first-fit?因为best-fit策略能显著降低后续大内存块的分配失败率。设想这样一个场景:系统里有两块空闲内存,一块是100MB,另一块是16MB,现在来了一个需要15MB的张量。如果用first-fit,可能随机分配了100MB那块,剩下85MB的碎片;而best-fit会选择16MB那块,100MB那块能完整保留给更“胖”的后续张量。实测下来,在典型的卷积网络里,best-fit比first-fit能多复用大概5%到8%的内存。

不过贪心并不保证全局最优。一个典型的反例是:两个大块内存的lifetime先结束,中间穿插着两个小块内存;贪心算法可能会把大块内存先复用小块的请求,后面真正需要大块的张量来了只能新开一块内存。如果改用回溯搜索,可能一开始就让小内存复用另一块,保留下完整的大块。好在这个问题在实际网络中出现概率不高,而且metadef支持一个“多次规划、取最优”的开关:可以连续尝试几种不同的启发式策略,取峰值内存最小的一次结果。

4. 实际内存规划过程与参数选择

4.1 峰值内存的评估公式

在编译阶段,metadef需要给上层下达一个指令:这一整张图,芯片至少要预留多少显存?这个数字通常不会是一个精确的固定值,而是一个上界评估。

假设图上所有不可复用内存块的大小总和为S_total,可复用张量经过复用规划后的新增内存总大小为S_reuse,算子本身的固定工作空间为S_workspace,那么峰值内存预估可以粗略写成:

Mem_peak = S_total + S_reuse + S_workspace

其中S_reuse的大小由上面的贪心复用算法决定。如果你关闭所有复用,那S_reuse就等于所有中间张量大小之和,峰值内存会高出非常多。

有一个经验值可以分享:对于一个标准的ResNet-50推理图,如果关闭内存复用,中间张量峰值大约在800MB左右;打开复用之后,这个值大概可以压到300MB以内。训练场景的差距更大,因为训练图里反向传播会额外产生大量激活张量,激活重计算开启与否,内存差两三倍都很正常。

4.2 对齐规则与内存偏移计算

内存规划不只是“大小匹配”,还有一个特别容易被忽略的维度——对齐。昇腾芯片上的内存操作通常要求对齐到一定的字节数,比如32字节或64字节。原因在于硬件搬运单元(类似DMA)按块读取数据,如果地址不对齐,轻则多读几轮、性能下降,重则直接触发地址错误。

在连续内存池模式下,每个张量的起始偏移需要经过对齐计算。如果一个张量的大小是4KB,但下一个张量要求64字节对齐,那么它的起始偏移必须是64的倍数。计算公式很简单:

offset_next = align_up(offset_cur + size_cur, alignment)

其中align_up的意思是把数值向上取整到alignment的整数倍。

这个对齐计算单独看没什么难度,但结合复用就有陷阱了:同一个物理内存块可能在生命周期不同的阶段被多个张量复用,而这些张量的对齐要求各不相同。规划时必须以“最严格的对齐要求”为准,否则某个张量会被分配到一个非法地址。我在看一些早期实现时遇到过这样一个问题——第一轮复用的张量是32字节对齐的,规划得非常舒服;第二轮复用一个64字节对齐的张量,结果它被分配到了32字节对齐的地址上,算子执行时直接报错。这个问题的排查很隐蔽,因为你看到的错误信息只会说“memory address error”,不会告诉你根因是对齐级别不够。

4.3 工作空间(Workspace)的复用策略

算子内部的工作空间,是内存规划里另一个大头。很多高性能算子内部需要一块临时buffer,大小可能和输入张量尺寸、卷积kernel尺寸、切分策略都相关。同一个算子在不同shape下的workspace大小可能差出数倍。

metadef对workspace的处理方式和普通中间张量略有不同。普通中间张量的生命周期完全由数据依赖决定,workspace的生命周期则是“从算子开始执行到算子结束执行”。理论上,如果两个算子在执行顺序上完全互斥(一个结束了另一个才开始),它们的workspace就可以复用同一块内存。这个推论实现起来很容易,但很多人忘了做。

实操中的做法是:把每个算子的workspace想象成一块“临时租赁地”,只在整个图的一个时间窗口内有效。然后在复用规划阶段,把workspace也加入全局内存池的待分配集合,和普通中间张量一视同仁。比如在做Transformer推理时,Attention块里的QKV投影算子和FFN里的第一个全连接算子,各自需要不小的workspace,但它们不会同时执行,完全可以让它们复用同一块区域,省下几百MB都很正常。

5. 常见问题与排查技巧实录

5.1 显存峰值异常偏高的五个场景

我在实际分析和排查问题上,把显存峰值异常的情况整理成了一个速查表,你可以直接对照着看。

场景表层现象深层原因解决方案
静态shape模型内存峰值翻倍编译通过,运行时申请内存过大复用时序被破坏,大部分张量生命周期重叠检查算子融合选项,减少中间张量输出
动态shape模型内存持续增长每次推理显存占用缓慢上涨动态shape导致内存池不断扩展创建新块,旧块无法复用开启内存池的“回收复用”选项,限制池内块数量
多路并发推理显示内存超卖并发路数增加后内存暴涨每路推理实例独立规划内存池,互不复用开启图级内存共享,多实例复用同一底层缓冲
单算子workspace优化后内存无变化算子内部buffer减小但总量不变workspace未参与全局复用规划让workspace申请信息上报到metadef统一规划
不同格式的同一张量占用双倍内存内存池里出现两个大小相同但生命周期重叠的块数据格式转换产生额外中间张量,且转换算子未被融合开启格式转换算子与相邻算子的融合优化

这个表格里的前两条是我见到最多的。尤其是动态shape场景,那不是复用协议本身有问题,而是复用决策依赖生命周期信息,而动态shape下生命周期信息在编译期永远无法完全确定。这个矛盾没有完美的解,只能靠运行时内存池管理来兜底。

5.2 一次真实的复用失效排查过程

有一回我在优化一个视觉模型的推理内存时,发现内存峰值比理论公式算出来的高得多。按照公式,压完复用之后S_reuse应该只有200MB,但实测峰值有500MB。我原以为是复用算法出了问题,后来一步一步查下去,才发现问题出在算子融合策略上。

那个模型里有两个相邻算子,理论上可以把其中一个算子的输出直接“透传”给下一个算子,不产生中间张量。但融合策略并没有把这两个算子合并,因为它们在GE层面属于不同的控制流分支,编译器无法在静态分析阶段确认分支会不会互相影响。于是每个分支都各自保留了自己的中间输出,内存占用自然降不下来。

这类问题的排查路径是:先看内存池的总大小,再看哪些节点占的体积最大,最后倒推这些节点为什么没有被合并掉。很多情况下,问题不在复用算法本身,而在上游图优化阶段已经埋下了隐患。这也印证了metadef设计的合理性——内存问题要从图构建阶段就开始控制,而不是等到规划阶段再亡羊补牢。

5.3 越界问题的定位顺口溜

内存复用带来的一个副产品是“越界很难查”。同一块物理内存在不同时间被多个张量使用,如果其中某个算子写越界了,它可能不会立刻报错,而是破坏了下一次复用的那个张量的数据。这导致的错误非常诡异:可能是某个输出张量数值偶尔不对,可能是随机NaN,也可能是完全没有规律的错误结果。

定位这类问题,我自己总结了一个顺口溜:先复现,再隔离;先查新算子,再查老算子;先关复用,再二分定位。具体操作是先把内存复用关掉,跑同一份输入数据,如果错误消失,就说明大概率是越界踩了别人的内存而不是自身逻辑错误。然后重新开启复用,每次只保留一个候选算子,用二分法确定是哪个算子的越界写坏了别人。

这个排查过程看着笨重,但确实是目前最可靠的手段。毕竟在复用协议生效的情况下,错误信息不会告诉你“谁的地址写爆了”,只会告诉你“数据对不上”。

5.4 内存碎片的“隐形消耗”

最后再说一个容易被忽略的现象:碎片。连续内存池模式在规划阶段已经把偏移排布定死了,不存在运行时碎片问题。但离散内存池模式不同,随着多次分配和释放,池里会出现大量不连续的小空闲块。这些小空闲块单独看都太小,无法分配给稍大的张量,但它们加起来可能占几百MB。

解决碎片问题的思路是在metadef层做一次“压缩重排”:在图的某个同步点,把所有存活张量按照旧的偏移拷贝到新的紧凑偏移上,然后把空闲块合并成更大的连续区域。这个操作在训练场景里是可以接受的,因为训练本身有同步点;但在推理场景里,重排操作会阻塞流水,代价比较大,通常建议在模型加载阶段就做一次静态压缩,运行阶段尽量不触发重排。

写在后面的一些实操体会

我实际使用和阅读metadef相关实现时,最深的感触是:内存复用协议的本质是“对时间的精确管理”,而不是“对空间的粗暴压缩”。空间复用只是结果,时间轴上的生命周期分析才是核心。谁把生命周期看得越准,谁就能在更小的显存上跑起更大的模型。

另外一个比较容易被忽视的经验是,不要迷信某个单一指标。只看“峰值内存”可能会让你在优化方向上跑偏。有一次我把峰值内存压了30%,结果算子执行时拷贝开销变大,整体端到端时延反而涨了15%。原因是我为了复用某块地址,让两个算子共用内存后引入了额外的数据拷贝等待。内存节省和性能之间永远是一笔需要综合评估的账。

如果你正在做类似的内存优化工作,建议先从可视化生命周期开始:把你图里的每个张量画成一条生命周期横线,找出那些“长尾”张量——它们往往是最值得优化的对象。很多长尾张量的生命周期之所以长,不是因为真的需要活那么久,而是因为图优化阶段少做了一次算子融合或者数据格式转换。修掉这些,内存复用协议才能真正发挥出它的威力。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询