想象你是一位舞台监督,正在为一台多幕舞台剧做准备。后台的储物仓很小,却要容纳几十件道具:开场要用的花篮、第二幕的书桌、第三幕的马车……你不可能为每件道具都准备一个专属柜子,好在很多道具永远不会同时登台。于是你翻开剧本,给每件道具标出“出场时间”和“退场时间”,让从不碰面的道具共用同一个柜子。这个“道具排班表”,就是内存规划。 AI 编译器面对的模型,就像这台舞台剧:每个算子的输出张量是一件道具,内存是后台仓库。内存规划(Memory Planning)为每一个张量安排内存位置,并在不需要时快速腾出空间。好的规划能让很有限的显存撑起很大的模型;糟糕的规划则会让 GPU 在某个瞬间爆显存。学会内存规划,等于掌握了“在狭小后台排演大戏”的舞台监督手艺。
对 AI 从业者来说,内存规划是理解显存占用、读懂编译器报错、在训练和推理中系统性降低峰值内存的必备技能。它把“爆显存怎么办”从一次次盲试,变成有章法的排查与优化。
可视化演示
展示内存规划三维度:张量生命周期分析(每个张量从产生到最后使用构成区间扫描线时刻存活张量内存之和为峰值)、冲突图着色分配(生命周期重叠的张量连冲突边着色算法同色节点共享内存槽最少颜色数为峰值)、峰值内存对比(无规划每个张量独立分配峰值=总和有规划重用内存槽峰值降至最大并发占用降低30%-70%),三模式切换理解AI编译器内存管理核心技术。
本节配有交互动效(进入视口后自动加载,约 3KB)。
'+C1+'
基本概念
舞台监督首先要回答一个最基础的问题:后台里需要管理的“道具”到底是什么?在 AI 编译器眼里,计算图中的每一个张量就是一件道具。输入数据是开场前就搬进后台的固定道具,最终输出是谢幕时才需要的压轴道具,而真正让内存吃紧的,是每一层算子之间产生的中间结果——它们像从舞台上匆匆走过的半成品道具,数量多、寿命短,却占据了大量后台空间。
这些中间结果也叫“中间缓冲”。它们每时每刻都在被创建和销毁,如果每个中间缓冲都独占一块内存,再大的显存也会被瞬间填满。内存规划要做的,就是把这些“半成品道具”放到最合适的柜子里,让它们在需要时出现、不需要时立刻腾位。显存管理因此从“每个张量问系统要内存”变成“由编译器统一安排床位”。
具体来说,编译器会先解析模型得到计算图,图中节点是算子,边是张量。内存规划会根据每个张量的大小、数据类型和执行顺序,确定它应该落在哪块地址上。一个基本的事实是:任何时刻所需的内存,等于该时刻所有活跃张量的大小之和。
所以,基本功就是认清张量的出生和死亡。只有弄清楚哪些张量会同时存在,才能知道哪些柜子可以共享。下一节我们就来看,舞台监督是怎么给道具登记“生卒时间”的。
M(t) 尽量小。- 识别张量集合
\mathcal{A}_t表示时刻t仍然存活的所有张量集合,也就是“这一刻还占用柜子的道具”。 - 单个张量大小
\mathrm{size}(a)表示张量a需要占用的内存字节数,类似于一件道具的占地面积。 - 累加求和
\sum把该时刻所有活跃张量的大小相加,得到当前瞬间的内存总需求。 - 每一刻的最小值
M(t)是时刻t的最低内存需求,编译器要尽力压低它的峰值。
生命周期分析
舞台监督排班的第一步,不是急着搬道具,而是翻开剧本,给每件道具标出两个时间:什么时候需要它上场,什么时候它再也不会被用到。这两个时间之间的区间,就是道具的“生命期”。张量也一样,生命周期是内存规划的地基。
一个张量的生命周期从产生它的算子执行完毕时开始,到它最后一次被某个算子消费时结束。前半段时间叫“出生时刻”,后半段叫“最后使用时刻”。编译器通过数据依赖分析,能精确推出每个张量被哪些算子使用,进而确定它是否还“活着”。如果某个张量出生后很快就没有消费者了,它的生命就结束了。
这张生卒时间表最终会形成每个时刻的“活跃集”。活跃集中的张量必须同时待在内存里,而不在活跃集中的张量,就是可以被复用的“死道具”。用区间表示一个张量的生命周期,是所有后续算法的起点。
有了这张生卒时间表,舞台监督就能轻松发现:A 道具在第三幕撤场,C 道具在第四幕才登场,它们永远不会在后台同时占地方。于是,内存复用算法就可以登场了。
t_last_use 之后,这件道具才真正被释放。- 定义张量
a代表一个具体的张量,比如某层卷积的输出。 - 出生时刻
t_{\mathrm{def}}(a)表示张量a被计算出来的时间点,也就是它的定义点。 - 最后使用时刻
t_{\mathrm{last\_use}}(a)表示张量a最后一次被消费的时间点。 - 区间表示
L(a)用闭区间把这两个端点括起来,代表张量的完整存活期。
内存复用算法
舞台监督最得意的手艺,就是让从不碰面的道具共用一个柜子。比如第三幕的花篮和第二幕的书桌,虽然都很大,但一个撤了另一个才上,于是它们共用同一个后台角落。在内存规划里,这叫“内存复用”。
实现这种复用最常见的方法是贪心算法。编译器把张量按照出生顺序排好队,然后逐个处理:轮到某个张量时,看现有内存槽里有没有在它出生前已经释放的槽,如果有就放进去;如果没有,就新开一个槽。这个过程很像酒店前台的房间分配:新客人来了,先查有没有已退房的空房间,没有才开新房间。
另一种思路是把问题建模成“区间图着色”。每个张量是一个区间,两个区间重叠意味着它们不能同色,不重叠则可以共用一种颜色。颜色总数就是需要的槽位数。图着色是经典 NP 问题,实践中常使用贪心近似,但能提供比直观方法更优的摆放方案。
贪心分配虽然简单,但在真实编译器中还有一个问题:如果内存槽的大小差别很大,仅仅“有空隙”还不够,必须“刚好放得下”。为了减少这种浪费,编译器需要把仓库本身也装修一番,这就引出了内存池。
a_i 挑一个在它出生前就已经空闲的最小编号槽 k。如果没有这样的槽,就新开一个。- 当前张量
a_i是正在被安排的第i个张量。 - 张量出生时刻
t_{\mathrm{def}}(a_i)表示该张量要入住内存的时间。 - 槽位空闲时间
t_{\mathrm{free}}(k)表示编号为k的内存槽上次释放的时间,只有当前张量出生时它已经空闲,才能复用。 - 取最小槽号
\min表示在所有满足条件的槽中选最靠前的那个,让内存布局更紧凑。
内存池机制
单纯地“共用一个柜子”还不够。如果每个张量都临时找系统要一小块内存,就会像演员每次上场前都去库房现租一个箱子,不仅慢,还会把后台堆得七零八落。所以舞台监督会在后台圈出一块专用区域,提前切成大小不一的格子,演出期间所有道具都在这里随取随放。这就是内存池。
内存池的基本思想是:在程序启动时,一次性从设备侧申请一大块连续显存,然后按常见的张量大小切分成多个块(chunk)。分配张量时,从池中取一个足够大的块;释放张量时,把块还回池中而不是还给操作系统。这样做最大的好处是减少了显存分配的系统调用,速度可以提升几个数量级,同时避免频繁向设备申请/释放造成的显存碎片。
碎片是内存池要治理的核心问题。如果反复分配和释放各种大小的块,空闲空间会被切得七零八落,导致明明总空闲内存够用,却找不到一块能满足请求的连续区域。衡量碎片程度可以看一个指标:
内存池让分配速度变快了,但碎片依然可能存在。要判断内存规划做得好不好,光看“有没有爆”还不够,还需要一套度量指标。如果说内存池是后台的装修,那么性能指标就是监督用来汇报工作的数字报表。
F 越接近 1,说明空闲内存被切得越碎;越接近 0,说明空闲空间越集中、越容易被利用。- 总空闲空间
\mathrm{size}(\mathrm{total\_free})表示当前内存池里所有空闲块大小的总和。 - 最大空闲块
\mathrm{size}(\mathrm{largest\_free})表示所有空闲块中最大的那一块的大小。 - 可用比例两者相除,得到当前空闲内存中“最大连续可用区域”的占比。
- 碎片程度用 1 减去这个比例,就是把碎片化程度归一化到 0 到 1 之间。
性能指标
舞台监督总结工作成绩时,不会只说“道具都放得下”,而是会报出几个硬数字:“整场演出后台最多同时堆了 12 件道具,占后台总面积的 60%,比原始方案省了 40% 的地方。”这几个数字,就是内存规划的性能指标。
最重要的指标是峰值内存,它是一台设备在运行过程中任意时刻最高需要占用的内存量。另一个直观指标是内存复用率:如果不做任何复用,把所有张量都独立分配,总内存需求记为 S_naive;经过规划后,实际峰值被压到 M_peak。两者差值占 S_naive 的比例,就是复用率。
此外,内存占用率描述的是峰值内存占设备总显存的比例;带宽利用则关注数据搬运量,因为内存规划不只影响“用多少地方”,还影响“搬多少数据”。好的规划应该让相邻计算尽量复用缓存中的数据,减少访存延迟。
不过,指标只是结果。要真正降低峰值内存,监督有时需要主动“改剧本”:让某件道具就地变装,或者把两幕后台合成一幕。这就是下一节要讨论的优化策略。
ρ 越接近 1,说明内存越被高效利用;如果为 0,说明规划前后没有任何区别。- 实际峰值内存
M_{\mathrm{peak}}表示规划后整个运行过程中的最大内存占用。 - 朴素总需求
S_{\mathrm{naive}}表示每个张量都独占内存时,所有张量大小之和。 - 节省比例两者相减后再除以
S_{\mathrm{naive}},得到被压缩掉的内存占比。 - 变成百分比
\rho就是复用率,用来横向比较不同规划方案的优劣。
优化策略
舞台监督不满足于“排好档期”,还会主动改剧本减少道具数量。比如一件道具下幕要变成另一件,那就让它直接原地改装,不需要再准备一件;如果两场戏中间没有换景时间,干脆把两场戏合成一场,中间的道具就不用搬出来。内存优化的三大策略,思路完全一样。
第一种是原地操作(In-place)。很多算子,比如 ReLU、某些归一化操作,输出和输入形状相同,且计算时不再需要原始输入,这时可以让输出直接覆盖到输入的内存上。这样两个张量只需要一份空间。第二种是算子融合:把多个连续算子合并成一个算子,中间结果只存在于寄存器或片上缓存中,完全不落到全局显存。第三种是内存交换:当峰值马上要来临时,把部分暂时不用的张量搬到 CPU 内存,等用的时候再搬回,用 PCIe 带宽换显存空间。
另外,训练场景里还有一个常用策略叫梯度节省,也叫激活重计算:反向传播原本需要保存前向的所有中间激活,如果内存不够,就删掉一部分,反向计算时再重新算一遍,用时间换空间。
这些策略听起来很美,但实际落地时还要考虑算子是否支持原地、反向传播是否依赖原始值、数据换出换入的耗时是否可接受。真正的功夫,藏在工程实践里。
x 和输出 y 共用一块内存,所以总占用只需要两者中较大的那一个,等于“一件道具改一改当两件用”。- 输入张量大小
\mathrm{size}(x)是原操作的输入张量占用的内存。 - 输出张量大小
\mathrm{size}(y)是原操作的输出张量占用的内存。 - 取最大值
\max表示因为输入输出共用空间,内存需求只需按最大的那个准备。 - 结果
M_{\mathrm{inplace}}是经过原地操作改造后的内存需求,通常远低于两个张量单独占用的总和。
工程实践
再好的排班方案,也要经过彩排验证。舞台监督的工作流程是:先通读剧本,再画时间表,接着分柜子,最后彩排测量,如果某一幕太挤,就回头调整。内存规划在工程中的落地,也遵循着同样的循环。
第一步是解析模型,把神经网络变成计算图;第二步是生命周期分析,标记每个张量的区间;第三步调用内存规划器,基于区间信息完成槽位分配;第四步生成带显存复用信息的执行代码;第五步用 profiler 工具测量峰值和耗时。如果结果不满足需求,就回到第三步,尝试更激进的复用或算子融合。这也是 PyTorch、MindSpore、TVM 等框架内部都在做的事。
以 PyTorch 为例,它默认使用 Caching Allocator 做设备端缓存分配,并且提供 torch.cuda.memory_allocated()、torch.cuda.max_memory_allocated() 等接口帮助你定位峰值。实测中,同样的 transformer 模型,不优化可能占用 11GB 显存,按生命周期做内存池规划后可能只需要 7GB。工程调优的经验法则只有一条:先测量,再优化,再测量。
从舞台监督到编译器工程师,内存规划始终围绕同一个道理:让有限的仓库,通过精细的排班和复用,支撑起更大的演出。看懂生命周期,用好复用与池化,再配上测量工具,你就能在显存这个“后台”里游刃有余。
O(n log n),即排序占主导。这意味着内存规划本身很快,不会成为编译期的瓶颈。- 张量数量
n表示计算图中参与规划的张量个数。 - 排序开销
\log n来自对张量按出生或死亡时间排序的复杂度。 - 单遍扫描排序后只需 O(n) 次扫描即可完成贪心分配,因此整体由排序主导。
- 表示成复杂度
O(n \log n)说明算法随模型规模增长是可接受的。
AI 应用场景
大模型训练显存优化
阶段:模型训练与系统调优
解决的问题:大模型训练时,激活值、梯度、优化器状态会瞬间挤爆显存。传统做法是盲目调小 batch size 或换更大的 GPU,AI 编译器用内存规划分析张量生命周期,自动复用中间缓冲,并配合激活重计算把峰值压下来。
具体案例:在 PyTorch 中训练 GPT 系列模型时,通过 torch.cuda.memory_allocated() 定位到峰值出现在 Transformer Block 的中间激活上。打开编译器级优化(如 TorchInductor)后,系统自动做算子融合和显存复用,同样的 13B 模型在 24GB 显卡上从“放不下”变成“放得下且可训练”。
推理服务的内存池化
阶段:模型部署与推理服务
解决的问题:在线推理要求低延迟,如果每次请求都动态分配显存,频繁调用会带来高延迟和碎片。内存池在服务启动时预分配好显存,按常见张量尺寸分块,请求之间复用块,减少分配开销。
具体案例:TensorRT 在构建推理引擎时,对每个卷积、Transformer 层的输出张量做生命周期分析,生成固定内存池。上线后服务 100 路并发时,显存占用比未优化的运行时低约 35%,P99 延迟下降 20%。
边缘设备上的内存规划
阶段:端侧模型编译与部署
解决的问题:手机、摄像头等边缘设备显存/缓存极其有限,还要求低功耗。传统直接运行深度学习框架开销太大;编译器先做内存规划,把中间结果尽量限制在片上缓存中,减少对 DRAM 的访问。
具体案例:TVM 在 ARM 设备上部署目标检测模型时,通过内存规划 + 算子融合,把模型运行峰值内存从 1.8GB 降到 620MB,首次加载时间缩短一半,手机端也能流畅跑实时检测。
Python 演示
class MemoryPlanner: def __init__(self): self.slot_end_list = [] def fit(self, tensors): order = sorted(range(len(tensors)), key=lambda i: tensors[i][1]) plans = [None] * len(tensors) for i in order: name, start, end, size = tensors[i] placed = False for s in range(len(self.slot_end_list)): if start >= self.slot_end_list[s]: plans[i] = s self.slot_end_list[s] = end placed = True break if not placed: plans[i] = len(self.slot_end_list) self.slot_end_list.append(end) return plans def calc_peak(tensors, plans): times = sorted(set([t[1] for t in tensors] + [t[2] for t in tensors])) peak = 0 for tt in times: used = sum(t[3] for t in tensors if t[1] <= tt < t[2]) peak = max(peak, used) return peak if __name__ == "__main__": tensors = [ ("A", 0, 2, 4), ("B", 1, 4, 6), ("C", 3, 6, 4), ("D", 5, 7, 8), ("E", 2, 3, 2), ("F", 4, 5, 2), ] planner = MemoryPlanner() plans = planner.fit(tensors) slots = max(plans) + 1 naive = sum(t[3] for t in tensors) peak = calc_peak(tensors, plans) rate = 1 - peak / naive print("张量槽位分配:", dict(zip([t[0] for t in tensors], plans))) print("内存槽数量:", slots) print("朴素峰值内存:", naive, "MB") print("规划后峰值内存:", peak, "MB") print("复用率: {:.2f}".format(rate))
self.slot_end_list = []记录每个内存槽上一次被释放的时间,初始为空。order = sorted(range(len(tensors)), key=lambda i: tensors[i][1])按张量出生时间排序,保证先出生的先分配。if start >= self.slot_end_list[s]检查槽s是否已经空闲,空闲才允许复用。plans[i] = s把当前张量放到这个空闲槽,记录分配结果。self.slot_end_list[s] = end更新该槽的最后释放时间,供后续张量判断。plans[i] = len(self.slot_end_list)没有空闲槽时,新开一个槽。calc_peak(tensors, plans)模拟所有关键时间点,计算任意时刻活跃张量大小总和的最大值。rate = 1 - peak / naive用朴素峰值与规划后峰值的差值计算复用率。
常见误区
误区我心里的内存管理是运行时的黑盒,框架会自动处理,出现显存不足只能换更大显卡或者调小 batch size,自己没法主动插手。
正解内存规划不是纯自动的魔法,它发生在编译器的中间表示阶段,工程师可以用内存池、生命周期分析、算子融合等手段主动干预。就像舞台监督能改道具排班表,而不是只能抱怨后台太小。
为何常见:日常使用 PyTorch 时,torch.cuda 的内存分配看起来是自动的,很少暴露出生命周期分析的过程,导致大家误以为无法优化。
误区我感觉神经网络里每个中间张量都会一直待在显存里,前向传播算完后,中间结果还在,所以显存占用会不断累积。
正解张量的生命周期到它最后一次被使用就结束了。编译器明确知道每个张量的死亡时刻,在后向传播开始前,大量中间激活就会被释放。内存规划正是利用“死亡”信息来复用这些空间。
为何常见:很多人在实验中只看到 torch.cuda.memory_allocated() 不断攀升,却忽略了某些阶段会自动释放,因为实际观测的粒度不够细。
误区只要用了内存池,显存碎片就为零了,内存池越大越好。
正解内存池只是把分配行为限制在池内,减少对操作系统的调用,但碎片依然可能出现。它更像是把零散的出租仓库整合成一栋大楼,若租户的占用时间乱,楼内依然会留下不好用的空房间。碎片率需要用公式衡量,并配合块大小对齐来治理。
为何常见:看到“预分配大块内存”的字面意思,容易以为所有地址都被牢牢控制,不会再产生浪费。
误区能复用就复用,最好把所有能共用的内存全部合并,峰值内存降到最低就是最优。
正解过度复用可能让两个算子争抢同一块缓存地址,导致访存冲突、流水线停顿。实际工程要在峰值、带宽、延迟和复杂度之间折中。复用率高是方向之一,但不是唯一指标,就像道具共柜省了空间,却可能让演员换场时跑得更远。
为何常见:初学者往往只盯住显存这一个瓶颈,忽略了内存规划本身也会影响计算访存局部性和编译时间。
面试题
Q1. 内存规划的核心目标是什么?
- A.尽量用满 GPU 显存
- B.降低峰值内存占用← 正确答案
- C.提升 GPU 的核心频率
- D.减少神经网络的层数
答案:B。内存规划的核心是让同时活跃的张量占用尽量少的内存,从而降低峰值,而不是把显存堆满。GPU 频率和网络层数都不是内存规划的目标。
Q2. 张量的生命周期指的是什么?
- A.从创建到最后使用的时间段← 正确答案
- B.从读取到写入的时间段
- C.网络训练所消耗的总时长
- D.显存条物理寿命的长短
答案:A。生命周期从张量的定义时刻开始,到最后一次被消费的时刻结束。它不关心读写之间的细微间隔,也不关心硬件寿命。
Q3. 有四个张量生命周期分别为 [0,3]、[2,5]、[4,7]、[6,8],贪心分配最少需要几个槽位?
- A.仅需要 1 个
- B.需要 2 个← 正确答案
- C.需要 3 个
- D.需要 4 个
答案:B。贪心分配:A 放槽 0,B 与 A 重叠放槽 1,C 在 A 结束后可放进槽 0,D 在 B 结束后可放进槽 1,所以总共只需要 2 个槽。
Q4. 内存池最主要的好处是?
- A.减少 GPU 显存分配的系统开销← 正确答案
- B.增加 GPU 物理显存总容量
- C.加快单个算子的数学计算速度
- D.缩小模型参数的存储体积
答案:A。内存池从设备侧一次性申请大块显存,在池内切块复用,从而减少运行时分配器的调用次数。它并不会增加物理显存,也不会加快算子本身的计算。
Q5. 下列哪种做法最能直接提升内存复用?
- A.把 ReLU 改为原地操作← 正确答案
- B.将训练的 batch size 调大
- C.使用更高精度的浮点数
- D.关闭框架的内存池机制
答案:A。原地操作能让输出覆盖输入,两个大小相同的张量只占一份内存,直接降低峰值。调大 batch size 反而增加内存,更高精度也会增大占用。
Q6. 影响内存规划效果的关键因素包括哪些?
- A.张量的生命周期信息← 正确答案
- B.算子的底层内核实现← 正确答案
- C.模型的参数总量大小
- D.算子的执行调度顺序← 正确答案
答案:ABD。生命周期和调度顺序直接决定哪些张量可复用;算子底层实现决定它是否支持原地操作或融合。模型参数量是固定的,不是决定规划效果好坏的变量。
Q7. 关于峰值内存,下列说法正确的是?
- A.等于所有张量的大小之和
- B.是活跃张量总和的最大值← 正确答案
- C.会随复用率提高而下降← 正确答案
- D.只取决于单个张量的尺寸
答案:BC。峰值内存按时间扫描所有时刻的活跃张量总和,取最大值,所以它不等于简单求所有张量之和,也不只取决于单个张量。复用率提高后,活动张量共享空间,峰值自然会下降。
Q8. 为什么不能把所有算子都改成原地操作?
- A.会破坏反向传播依赖和输入复用← 正确答案
- B.会明显降低 GPU 工作频率
- C.会增加内存碎片的比例
- D.会导致显存永久泄漏
答案:A。很多算子在被原地覆盖后,原始输入会丢失。若反向传播还需要原始输入计算梯度,就无法原地操作。GPU 频率、碎片和泄漏都不是这一策略的主要限制。
Q9. 以下哪些是降低训练峰值内存的常见策略?
- A.使用梯度检查点技术← 正确答案
- B.融合多个相邻算子← 正确答案
- C.清空所有 GPU 显存缓存
- D.将张量换出到 CPU 内存← 正确答案
答案:ABD。梯度检查点用重计算换内存;算子融合避免中间结果落显存;内存交换把暂时不用的张量放到 CPU。清空显存缓存只是释放已缓存空间,不能从算法上降低峰值。
Q10. 内存规划通常发生在编译器流程的哪个阶段?
- A.算子调度完成之后← 正确答案
- B.模型训练收敛之后
- C.数据加载开始之前
- D.用户交互发生之前
答案:A。编译器先完成计算图解析和算子调度,才能知道张量的生命周期和执行顺序,因此内存规划发生在调度之后、代码生成之前。它与训练收敛、数据加载、用户交互无关。