AI 编译器内存分配算法:从静态/动态划分到内存复用

AI 编译器内存分配算法:从静态/动态划分到内存复用

训练一个模型,显存到底花在哪了?

答案不是"权重太大",而是激活值、梯度、优化器状态 这几座大山。AI 编译器的内存分配,做的就是在正确的时间,把同一块显存借给不同的人用。

这篇基于 AISystem 前端优化章节,并修正了原文自己标注为"不太准确"的静态/动态内存划分。

文章目录

  • [AI 编译器内存分配算法:从静态/动态划分到内存复用](#AI 编译器内存分配算法:从静态/动态划分到内存复用)
    • 一、为什么内存分配是编译器的活
      • [1.1 一个扎心的现实:算力涨得比显存快](#1.1 一个扎心的现实:算力涨得比显存快)
      • [1.2 训练时显存都被谁吃了](#1.2 训练时显存都被谁吃了)
      • [1.3 前端做决策,后端做执行](#1.3 前端做决策,后端做执行)
    • 二、内存怎么划分:修正原文的分类
      • [2.1 原文的划分,以及它自己指出的问题](#2.1 原文的划分,以及它自己指出的问题)
      • [2.2 更准确的划分:按"生命周期"分](#2.2 更准确的划分:按"生命周期"分)
      • [2.3 顺便回答原文的问题:N 层网络要多少显存](#2.3 顺便回答原文的问题:N 层网络要多少显存)
    • 三、内存复用:两种基本操作
      • [3.1 Inplace Operation(原地操作)](#3.1 Inplace Operation(原地操作))
      • [3.2 Memory Sharing(内存共享)](#3.2 Memory Sharing(内存共享))
    • [四、标签 + 内存池:一个可执行的分配算法](#四、标签 + 内存池:一个可执行的分配算法)
      • [4.1 算法流程](#4.1 算法流程)
      • [4.2 为什么这个算法像寄存器分配](#4.2 为什么这个算法像寄存器分配)
      • [4.3 内存池:别每次都 malloc](#4.3 内存池:别每次都 malloc)
    • 五、并行分配:别让"省内存"把并行性省没了
      • [5.1 两种方案,同样的串行语义](#5.1 两种方案,同样的串行语义)
      • [5.2 实际怎么权衡](#5.2 实际怎么权衡)
    • [六、MobileNetV2 的例子:优化前后差多少](#六、MobileNetV2 的例子:优化前后差多少)
    • 七、小结
    • 参考资料

一、为什么内存分配是编译器的活

1.1 一个扎心的现实:算力涨得比显存快

GPU CUDA 核心 FP32 算力 显存
Tesla K80 2496 8.7 TFLOPS 12 GB
Tesla P40 3840 12 TFLOPS 24 GB
Tesla V100 5120 14 TFLOPS 16 GB
Tesla T4 2560 8.1 TFLOPS 16 GB
Tesla V100S 5120 16.35 TFLOPS 32 GB
Tesla A100 6912 19.49 TFLOPS 80 GB

看懂这张表 :从 K80 到 A100,算力涨了约 2.2 倍 (FP32),显存涨了约 6.7 倍 ------但模型规模的涨幅远超这两个数 。所以真正的矛盾不是"显存不够大",而是显存增长速度追不上模型胃口。

勘误提示 :原文这张表的"CUDA 核心数量"一行有 7 个数字对 6 个列 (多出的 14592 无处安放),本文已按公开规格修正为 6 个值;另外 A100 的 INT8 并非 N/A(312 TOPS 稠密 / 624 TOPS 稀疏),功耗也分 SXM 400W 与 PCIe 250W。

1.2 训练时显存都被谁吃了

一个经典例子(原文数据):

模型 Batch Size 显存需求
Inception v4(ImageNet) 32 约 40 GB
Wide ResNet-152(ImageNet) 64 约 180 GB

参数量才多少? Inception v4 约 4300 万参数(FP32 才 170 MB)。40 GB 里绝大部分不是权重。
#mermaid-svg-nZYyxhO33fyDabis{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-nZYyxhO33fyDabis .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-nZYyxhO33fyDabis .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-nZYyxhO33fyDabis .error-icon{fill:#552222;}#mermaid-svg-nZYyxhO33fyDabis .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-nZYyxhO33fyDabis .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-nZYyxhO33fyDabis .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-nZYyxhO33fyDabis .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-nZYyxhO33fyDabis .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-nZYyxhO33fyDabis .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-nZYyxhO33fyDabis .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-nZYyxhO33fyDabis .marker{fill:#333333;stroke:#333333;}#mermaid-svg-nZYyxhO33fyDabis .marker.cross{stroke:#333333;}#mermaid-svg-nZYyxhO33fyDabis svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-nZYyxhO33fyDabis p{margin:0;}#mermaid-svg-nZYyxhO33fyDabis .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-nZYyxhO33fyDabis .cluster-label text{fill:#333;}#mermaid-svg-nZYyxhO33fyDabis .cluster-label span{color:#333;}#mermaid-svg-nZYyxhO33fyDabis .cluster-label span p{background-color:transparent;}#mermaid-svg-nZYyxhO33fyDabis .label text,#mermaid-svg-nZYyxhO33fyDabis span{fill:#333;color:#333;}#mermaid-svg-nZYyxhO33fyDabis .node rect,#mermaid-svg-nZYyxhO33fyDabis .node circle,#mermaid-svg-nZYyxhO33fyDabis .node ellipse,#mermaid-svg-nZYyxhO33fyDabis .node polygon,#mermaid-svg-nZYyxhO33fyDabis .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-nZYyxhO33fyDabis .rough-node .label text,#mermaid-svg-nZYyxhO33fyDabis .node .label text,#mermaid-svg-nZYyxhO33fyDabis .image-shape .label,#mermaid-svg-nZYyxhO33fyDabis .icon-shape .label{text-anchor:middle;}#mermaid-svg-nZYyxhO33fyDabis .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-nZYyxhO33fyDabis .rough-node .label,#mermaid-svg-nZYyxhO33fyDabis .node .label,#mermaid-svg-nZYyxhO33fyDabis .image-shape .label,#mermaid-svg-nZYyxhO33fyDabis .icon-shape .label{text-align:center;}#mermaid-svg-nZYyxhO33fyDabis .node.clickable{cursor:pointer;}#mermaid-svg-nZYyxhO33fyDabis .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-nZYyxhO33fyDabis .arrowheadPath{fill:#333333;}#mermaid-svg-nZYyxhO33fyDabis .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-nZYyxhO33fyDabis .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-nZYyxhO33fyDabis .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-nZYyxhO33fyDabis .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-nZYyxhO33fyDabis .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-nZYyxhO33fyDabis .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-nZYyxhO33fyDabis .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-nZYyxhO33fyDabis .cluster text{fill:#333;}#mermaid-svg-nZYyxhO33fyDabis .cluster span{color:#333;}#mermaid-svg-nZYyxhO33fyDabis div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-nZYyxhO33fyDabis .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-nZYyxhO33fyDabis rect.text{fill:none;stroke-width:0;}#mermaid-svg-nZYyxhO33fyDabis .icon-shape,#mermaid-svg-nZYyxhO33fyDabis .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-nZYyxhO33fyDabis .icon-shape p,#mermaid-svg-nZYyxhO33fyDabis .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-nZYyxhO33fyDabis .icon-shape .label rect,#mermaid-svg-nZYyxhO33fyDabis .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-nZYyxhO33fyDabis .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-nZYyxhO33fyDabis .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-nZYyxhO33fyDabis :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 训练显存
持久部分

(每个 step 都在)
瞬态部分

(随生命周期复用)
权重 θ
梯度 ∂L/∂θ
优化器状态

动量 m / 方差 v
常量(冻结权重等)
激活值

(前向输出,反向要用)
算子 workspace

(卷积算法临时区)
临时中间结果

这就是内存分配的价值 :持久部分省不了,瞬态部分可以借来借去。而瞬态部分恰恰是大头。

1.3 前端做决策,后端做执行

原文有一句很重要的定位:

内存分配"在前端优化中做分析和管理",但实际的分配动作在后端完成。

阶段 做什么
前端(图优化) 分析生命周期 、决定谁和谁可以共用一块内存、给出分配方案
后端(代码生成/运行时) 真正申请内存、生成地址、kernel 里用这些地址

类比 :前端是"排座位 ",后端是"摆椅子"。排座位的人不需要知道椅子是塑料的还是铁的,但必须知道谁什么时候来、什么时候走。


二、内存怎么划分:修正原文的分类

2.1 原文的划分,以及它自己指出的问题

原文把内存分成:

类别 原文包含
静态内存 模型权重、值节点(Value Node)、输出(Output)
动态内存 输出张量、工作区张量(Workspace)

但原文在正文里留了两段自我批注:

"上面的理解不太准确,请参考这个系列视频的理解,output 肯定是不对的 。值节点也不知道是什么,优化器参数和动量这个在哪里?"
"上面的理解不太准确,请参考这个系列视频的理解"

这种"作者自己标 bug"的写法很少见,但问题确实是真问题:

原文说法 问题
**输出(Output)**属于静态内存 输出是每次请求/每步训练产生的数据 ,不是"初始化时申请一次就不变"的东西;而且它和"输出张量"在同一份分类里重复出现
值节点(Value Node) 概念模糊,通常指图中的常量节点;可折叠的会被常量折叠干掉,不可折叠的才是常驻
优化器参数和动量 完全没提------而这恰恰是 Adam 训练里的大头(2 份与参数同大的状态)

2.2 更准确的划分:按"生命周期"分

与其纠结"静态/动态"这两个词,不如直接问:这块内存活多久?

类别 内容 生命周期
持久(persistent) 权重 θ 整个训练/推理
梯度 ∂L/∂θ 每步重算,但缓冲区常驻
优化器状态(SGD 动量;Adam 的 m、v;混合精度 master weights) 整个训练
不可折叠的常量(冻结权重、常量节点) 整个运行期
瞬态(transient) 激活值(前向各层输出,反向要用) 从产生到它的梯度算完
算子 workspace(卷积算法、reduce 临时区) 单次算子执行内
临时中间结果(未融合算子之间) 极短
外部持有 模型输入 / 模型输出 由调用方管理,不属于图内分配

#mermaid-svg-dunB4BWnmMf5Z68n{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-dunB4BWnmMf5Z68n .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-dunB4BWnmMf5Z68n .error-icon{fill:#552222;}#mermaid-svg-dunB4BWnmMf5Z68n .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-dunB4BWnmMf5Z68n .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-dunB4BWnmMf5Z68n .marker{fill:#333333;stroke:#333333;}#mermaid-svg-dunB4BWnmMf5Z68n .marker.cross{stroke:#333333;}#mermaid-svg-dunB4BWnmMf5Z68n svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-dunB4BWnmMf5Z68n p{margin:0;}#mermaid-svg-dunB4BWnmMf5Z68n .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-dunB4BWnmMf5Z68n .cluster-label text{fill:#333;}#mermaid-svg-dunB4BWnmMf5Z68n .cluster-label span{color:#333;}#mermaid-svg-dunB4BWnmMf5Z68n .cluster-label span p{background-color:transparent;}#mermaid-svg-dunB4BWnmMf5Z68n .label text,#mermaid-svg-dunB4BWnmMf5Z68n span{fill:#333;color:#333;}#mermaid-svg-dunB4BWnmMf5Z68n .node rect,#mermaid-svg-dunB4BWnmMf5Z68n .node circle,#mermaid-svg-dunB4BWnmMf5Z68n .node ellipse,#mermaid-svg-dunB4BWnmMf5Z68n .node polygon,#mermaid-svg-dunB4BWnmMf5Z68n .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-dunB4BWnmMf5Z68n .rough-node .label text,#mermaid-svg-dunB4BWnmMf5Z68n .node .label text,#mermaid-svg-dunB4BWnmMf5Z68n .image-shape .label,#mermaid-svg-dunB4BWnmMf5Z68n .icon-shape .label{text-anchor:middle;}#mermaid-svg-dunB4BWnmMf5Z68n .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-dunB4BWnmMf5Z68n .rough-node .label,#mermaid-svg-dunB4BWnmMf5Z68n .node .label,#mermaid-svg-dunB4BWnmMf5Z68n .image-shape .label,#mermaid-svg-dunB4BWnmMf5Z68n .icon-shape .label{text-align:center;}#mermaid-svg-dunB4BWnmMf5Z68n .node.clickable{cursor:pointer;}#mermaid-svg-dunB4BWnmMf5Z68n .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-dunB4BWnmMf5Z68n .arrowheadPath{fill:#333333;}#mermaid-svg-dunB4BWnmMf5Z68n .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-dunB4BWnmMf5Z68n .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-dunB4BWnmMf5Z68n .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-dunB4BWnmMf5Z68n .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-dunB4BWnmMf5Z68n .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-dunB4BWnmMf5Z68n .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-dunB4BWnmMf5Z68n .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-dunB4BWnmMf5Z68n .cluster text{fill:#333;}#mermaid-svg-dunB4BWnmMf5Z68n .cluster span{color:#333;}#mermaid-svg-dunB4BWnmMf5Z68n div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-dunB4BWnmMf5Z68n .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-dunB4BWnmMf5Z68n rect.text{fill:none;stroke-width:0;}#mermaid-svg-dunB4BWnmMf5Z68n .icon-shape,#mermaid-svg-dunB4BWnmMf5Z68n .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-dunB4BWnmMf5Z68n .icon-shape p,#mermaid-svg-dunB4BWnmMf5Z68n .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-dunB4BWnmMf5Z68n .icon-shape .label rect,#mermaid-svg-dunB4BWnmMf5Z68n .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-dunB4BWnmMf5Z68n .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-dunB4BWnmMf5Z68n .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-dunB4BWnmMf5Z68n :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 外部
输入 buffer
输出 buffer
瞬态
激活值
梯度算完即释放
workspace
算子结束即释放
持久
权重
梯度
优化器状态
常量

一句话记住 :"初始化时申请、全程不释放"的是持久内存;"用完就能借给别人"的是瞬态内存。 输出张量属于后者(而且更偏"外部持有")。

2.3 顺便回答原文的问题:N 层网络要多少显存

原文留了一个思考题:N 层网络需要多少 NPU 内存,怎么估?

推理(FP16):

M infer ≈ 2 P ⏟ 权重 + M workspace ⏟ 算子临时区 + M KV ⏟ LLM 才有 M_{\text{infer}} \approx \underbrace{2P}{\text{权重}} + \underbrace{M{\text{workspace}}}{\text{算子临时区}} + \underbrace{M{\text{KV}}}_{\text{LLM 才有}} Minfer≈权重 2P+算子临时区 Mworkspace+LLM 才有 MKV

训练(混合精度 + Adam) ,业界常用的粗略系数是 16P 字节:

项 字节数
FP16 权重 2P
FP16 梯度 2P
FP32 主权重(master weights) 4P
FP32 一阶动量 m 4P
FP32 二阶动量 v 4P
小计 16P

再加激活值:

M train ≈ 16 P + ∑ l = 1 N B ⋅ S ⋅ d ⋅ k l ⋅ bytes + M workspace M_{\text{train}} \approx 16P + \sum_{l=1}^{N} B \cdot S \cdot d \cdot k_l \cdot \text{bytes} + M_{\text{workspace}} Mtrain≈16P+l=1∑NB⋅S⋅d⋅kl⋅bytes+Mworkspace

其中 B B B 是 batch、 S S S 是序列长度、 d d d 是 hidden size, k l k_l kl 是该层每个 token 需要保存的中间张量个数(Transformer 常取 10~20 量级,取决于是否重计算)。

算个例子 :7B 模型、FP16 推理,权重约 14 GB ;训练用 Adam 混合精度,光"16P"就是 112 GB ------这才是单卡训不动的原因,跟激活值关系还不大。

所以:推理优化看权重 + KV Cache,训练优化先看优化器状态(ZeRO/FSDP 切的就是这 16P),再看激活值(重计算/梯度检查点)。


三、内存复用:两种基本操作

原文把节省内存的算法分成四类:

类型 做法 典型场景
空间换内存 把部分数据/模块卸载到 CPU 内存(CPU Offload) MoE 大模型、显存不够时
计算换内存 不存中间结果,需要时重算(Gradient Checkpointing) 训练深网络
模型压缩 量化(Quantization)、剪枝、蒸馏 端侧推理
内存复用 分析数据流,让不同张量共用一块内存 编译器主战场

勘误 :原文写"量化训练 Quantification "------正确拼写是 Quantization。

另外,"以空间换内存"这个说法有点绕:本质是用另一种存储(CPU 内存/磁盘)换显存 ,叫"offload"更准确。

这一节的重点是最后一类------内存复用。它有两个基本动作。

3.1 Inplace Operation(原地操作)

条件 :一块内存不再被需要,且下一个操作是 element-wise 的 → 可以直接在原地址上覆盖写。
#mermaid-svg-xKetU7fEbjnPxcbU{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-xKetU7fEbjnPxcbU .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-xKetU7fEbjnPxcbU .error-icon{fill:#552222;}#mermaid-svg-xKetU7fEbjnPxcbU .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-xKetU7fEbjnPxcbU .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-xKetU7fEbjnPxcbU .marker{fill:#333333;stroke:#333333;}#mermaid-svg-xKetU7fEbjnPxcbU .marker.cross{stroke:#333333;}#mermaid-svg-xKetU7fEbjnPxcbU svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-xKetU7fEbjnPxcbU p{margin:0;}#mermaid-svg-xKetU7fEbjnPxcbU .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-xKetU7fEbjnPxcbU .cluster-label text{fill:#333;}#mermaid-svg-xKetU7fEbjnPxcbU .cluster-label span{color:#333;}#mermaid-svg-xKetU7fEbjnPxcbU .cluster-label span p{background-color:transparent;}#mermaid-svg-xKetU7fEbjnPxcbU .label text,#mermaid-svg-xKetU7fEbjnPxcbU span{fill:#333;color:#333;}#mermaid-svg-xKetU7fEbjnPxcbU .node rect,#mermaid-svg-xKetU7fEbjnPxcbU .node circle,#mermaid-svg-xKetU7fEbjnPxcbU .node ellipse,#mermaid-svg-xKetU7fEbjnPxcbU .node polygon,#mermaid-svg-xKetU7fEbjnPxcbU .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-xKetU7fEbjnPxcbU .rough-node .label text,#mermaid-svg-xKetU7fEbjnPxcbU .node .label text,#mermaid-svg-xKetU7fEbjnPxcbU .image-shape .label,#mermaid-svg-xKetU7fEbjnPxcbU .icon-shape .label{text-anchor:middle;}#mermaid-svg-xKetU7fEbjnPxcbU .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-xKetU7fEbjnPxcbU .rough-node .label,#mermaid-svg-xKetU7fEbjnPxcbU .node .label,#mermaid-svg-xKetU7fEbjnPxcbU .image-shape .label,#mermaid-svg-xKetU7fEbjnPxcbU .icon-shape .label{text-align:center;}#mermaid-svg-xKetU7fEbjnPxcbU .node.clickable{cursor:pointer;}#mermaid-svg-xKetU7fEbjnPxcbU .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-xKetU7fEbjnPxcbU .arrowheadPath{fill:#333333;}#mermaid-svg-xKetU7fEbjnPxcbU .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-xKetU7fEbjnPxcbU .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-xKetU7fEbjnPxcbU .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xKetU7fEbjnPxcbU .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-xKetU7fEbjnPxcbU .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xKetU7fEbjnPxcbU .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-xKetU7fEbjnPxcbU .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-xKetU7fEbjnPxcbU .cluster text{fill:#333;}#mermaid-svg-xKetU7fEbjnPxcbU .cluster span{color:#333;}#mermaid-svg-xKetU7fEbjnPxcbU div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-xKetU7fEbjnPxcbU .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-xKetU7fEbjnPxcbU rect.text{fill:none;stroke-width:0;}#mermaid-svg-xKetU7fEbjnPxcbU .icon-shape,#mermaid-svg-xKetU7fEbjnPxcbU .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xKetU7fEbjnPxcbU .icon-shape p,#mermaid-svg-xKetU7fEbjnPxcbU .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-xKetU7fEbjnPxcbU .icon-shape .label rect,#mermaid-svg-xKetU7fEbjnPxcbU .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xKetU7fEbjnPxcbU .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-xKetU7fEbjnPxcbU .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-xKetU7fEbjnPxcbU :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 直接写回 x 的地址
输入 x
算子(element-wise)
输出 y

原文的例子(左图可行):
#mermaid-svg-4SdeeQkSjDcCA4bi{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-4SdeeQkSjDcCA4bi .error-icon{fill:#552222;}#mermaid-svg-4SdeeQkSjDcCA4bi .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-4SdeeQkSjDcCA4bi .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-4SdeeQkSjDcCA4bi .marker{fill:#333333;stroke:#333333;}#mermaid-svg-4SdeeQkSjDcCA4bi .marker.cross{stroke:#333333;}#mermaid-svg-4SdeeQkSjDcCA4bi svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-4SdeeQkSjDcCA4bi p{margin:0;}#mermaid-svg-4SdeeQkSjDcCA4bi .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster-label text{fill:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster-label span{color:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster-label span p{background-color:transparent;}#mermaid-svg-4SdeeQkSjDcCA4bi .label text,#mermaid-svg-4SdeeQkSjDcCA4bi span{fill:#333;color:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi .node rect,#mermaid-svg-4SdeeQkSjDcCA4bi .node circle,#mermaid-svg-4SdeeQkSjDcCA4bi .node ellipse,#mermaid-svg-4SdeeQkSjDcCA4bi .node polygon,#mermaid-svg-4SdeeQkSjDcCA4bi .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-4SdeeQkSjDcCA4bi .rough-node .label text,#mermaid-svg-4SdeeQkSjDcCA4bi .node .label text,#mermaid-svg-4SdeeQkSjDcCA4bi .image-shape .label,#mermaid-svg-4SdeeQkSjDcCA4bi .icon-shape .label{text-anchor:middle;}#mermaid-svg-4SdeeQkSjDcCA4bi .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-4SdeeQkSjDcCA4bi .rough-node .label,#mermaid-svg-4SdeeQkSjDcCA4bi .node .label,#mermaid-svg-4SdeeQkSjDcCA4bi .image-shape .label,#mermaid-svg-4SdeeQkSjDcCA4bi .icon-shape .label{text-align:center;}#mermaid-svg-4SdeeQkSjDcCA4bi .node.clickable{cursor:pointer;}#mermaid-svg-4SdeeQkSjDcCA4bi .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-4SdeeQkSjDcCA4bi .arrowheadPath{fill:#333333;}#mermaid-svg-4SdeeQkSjDcCA4bi .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-4SdeeQkSjDcCA4bi .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-4SdeeQkSjDcCA4bi .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-4SdeeQkSjDcCA4bi .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-4SdeeQkSjDcCA4bi .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-4SdeeQkSjDcCA4bi .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster text{fill:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi .cluster span{color:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-4SdeeQkSjDcCA4bi .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-4SdeeQkSjDcCA4bi rect.text{fill:none;stroke-width:0;}#mermaid-svg-4SdeeQkSjDcCA4bi .icon-shape,#mermaid-svg-4SdeeQkSjDcCA4bi .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-4SdeeQkSjDcCA4bi .icon-shape p,#mermaid-svg-4SdeeQkSjDcCA4bi .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-4SdeeQkSjDcCA4bi .icon-shape .label rect,#mermaid-svg-4SdeeQkSjDcCA4bi .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-4SdeeQkSjDcCA4bi .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-4SdeeQkSjDcCA4bi .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-4SdeeQkSjDcCA4bi :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} B 计算完
C 复用 B 的内存
E 复用 C 的内存

右图不可行:
#mermaid-svg-Qtz20n8G1SGdGaqm{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-Qtz20n8G1SGdGaqm .error-icon{fill:#552222;}#mermaid-svg-Qtz20n8G1SGdGaqm .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-Qtz20n8G1SGdGaqm .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-Qtz20n8G1SGdGaqm .marker{fill:#333333;stroke:#333333;}#mermaid-svg-Qtz20n8G1SGdGaqm .marker.cross{stroke:#333333;}#mermaid-svg-Qtz20n8G1SGdGaqm svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-Qtz20n8G1SGdGaqm p{margin:0;}#mermaid-svg-Qtz20n8G1SGdGaqm .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster-label text{fill:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster-label span{color:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster-label span p{background-color:transparent;}#mermaid-svg-Qtz20n8G1SGdGaqm .label text,#mermaid-svg-Qtz20n8G1SGdGaqm span{fill:#333;color:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm .node rect,#mermaid-svg-Qtz20n8G1SGdGaqm .node circle,#mermaid-svg-Qtz20n8G1SGdGaqm .node ellipse,#mermaid-svg-Qtz20n8G1SGdGaqm .node polygon,#mermaid-svg-Qtz20n8G1SGdGaqm .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-Qtz20n8G1SGdGaqm .rough-node .label text,#mermaid-svg-Qtz20n8G1SGdGaqm .node .label text,#mermaid-svg-Qtz20n8G1SGdGaqm .image-shape .label,#mermaid-svg-Qtz20n8G1SGdGaqm .icon-shape .label{text-anchor:middle;}#mermaid-svg-Qtz20n8G1SGdGaqm .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-Qtz20n8G1SGdGaqm .rough-node .label,#mermaid-svg-Qtz20n8G1SGdGaqm .node .label,#mermaid-svg-Qtz20n8G1SGdGaqm .image-shape .label,#mermaid-svg-Qtz20n8G1SGdGaqm .icon-shape .label{text-align:center;}#mermaid-svg-Qtz20n8G1SGdGaqm .node.clickable{cursor:pointer;}#mermaid-svg-Qtz20n8G1SGdGaqm .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-Qtz20n8G1SGdGaqm .arrowheadPath{fill:#333333;}#mermaid-svg-Qtz20n8G1SGdGaqm .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-Qtz20n8G1SGdGaqm .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-Qtz20n8G1SGdGaqm .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Qtz20n8G1SGdGaqm .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-Qtz20n8G1SGdGaqm .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Qtz20n8G1SGdGaqm .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster text{fill:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm .cluster span{color:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-Qtz20n8G1SGdGaqm .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-Qtz20n8G1SGdGaqm rect.text{fill:none;stroke-width:0;}#mermaid-svg-Qtz20n8G1SGdGaqm .icon-shape,#mermaid-svg-Qtz20n8G1SGdGaqm .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Qtz20n8G1SGdGaqm .icon-shape p,#mermaid-svg-Qtz20n8G1SGdGaqm .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-Qtz20n8G1SGdGaqm .icon-shape .label rect,#mermaid-svg-Qtz20n8G1SGdGaqm .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Qtz20n8G1SGdGaqm .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-Qtz20n8G1SGdGaqm .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-Qtz20n8G1SGdGaqm :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} B 计算完
C 想复用 B 的内存
但 F 还要用 B 的结果
❌ 不能复用

判据 :B 的结果在 C 之后还有没有别的消费者? 有就不能覆盖。

Inplace 的安全条件(工程上要同时满足):

条件 原因
算子是 element-wise 输出第 i 个元素只依赖输入第 i 个元素,不会读写冲突
shape / dtype 一致 缓冲区大小和解释方式要匹配
没有其他消费者 否则会读到被覆盖的值
不是图的外部输出 用户还要读它
反向不需要它的原始值 否则 autograd 会算错(这也是 PyTorch 里 inplace 常报 "a variable needed for gradient computation has been modified" 的原因)

3.2 Memory Sharing(内存共享)

条件 :两块数据大小相同 ,前一块已经死亡 (生命周期结束)→ 后一块可以直接接管它的内存,不需要重新申请/释放。
#mermaid-svg-MTAsSupYDSaVkI7O{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-MTAsSupYDSaVkI7O .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-MTAsSupYDSaVkI7O .error-icon{fill:#552222;}#mermaid-svg-MTAsSupYDSaVkI7O .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-MTAsSupYDSaVkI7O .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-MTAsSupYDSaVkI7O .marker{fill:#333333;stroke:#333333;}#mermaid-svg-MTAsSupYDSaVkI7O .marker.cross{stroke:#333333;}#mermaid-svg-MTAsSupYDSaVkI7O svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-MTAsSupYDSaVkI7O p{margin:0;}#mermaid-svg-MTAsSupYDSaVkI7O .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-MTAsSupYDSaVkI7O .cluster-label text{fill:#333;}#mermaid-svg-MTAsSupYDSaVkI7O .cluster-label span{color:#333;}#mermaid-svg-MTAsSupYDSaVkI7O .cluster-label span p{background-color:transparent;}#mermaid-svg-MTAsSupYDSaVkI7O .label text,#mermaid-svg-MTAsSupYDSaVkI7O span{fill:#333;color:#333;}#mermaid-svg-MTAsSupYDSaVkI7O .node rect,#mermaid-svg-MTAsSupYDSaVkI7O .node circle,#mermaid-svg-MTAsSupYDSaVkI7O .node ellipse,#mermaid-svg-MTAsSupYDSaVkI7O .node polygon,#mermaid-svg-MTAsSupYDSaVkI7O .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-MTAsSupYDSaVkI7O .rough-node .label text,#mermaid-svg-MTAsSupYDSaVkI7O .node .label text,#mermaid-svg-MTAsSupYDSaVkI7O .image-shape .label,#mermaid-svg-MTAsSupYDSaVkI7O .icon-shape .label{text-anchor:middle;}#mermaid-svg-MTAsSupYDSaVkI7O .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-MTAsSupYDSaVkI7O .rough-node .label,#mermaid-svg-MTAsSupYDSaVkI7O .node .label,#mermaid-svg-MTAsSupYDSaVkI7O .image-shape .label,#mermaid-svg-MTAsSupYDSaVkI7O .icon-shape .label{text-align:center;}#mermaid-svg-MTAsSupYDSaVkI7O .node.clickable{cursor:pointer;}#mermaid-svg-MTAsSupYDSaVkI7O .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-MTAsSupYDSaVkI7O .arrowheadPath{fill:#333333;}#mermaid-svg-MTAsSupYDSaVkI7O .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-MTAsSupYDSaVkI7O .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-MTAsSupYDSaVkI7O .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-MTAsSupYDSaVkI7O .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-MTAsSupYDSaVkI7O .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-MTAsSupYDSaVkI7O .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-MTAsSupYDSaVkI7O .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-MTAsSupYDSaVkI7O .cluster text{fill:#333;}#mermaid-svg-MTAsSupYDSaVkI7O .cluster span{color:#333;}#mermaid-svg-MTAsSupYDSaVkI7O div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-MTAsSupYDSaVkI7O .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-MTAsSupYDSaVkI7O rect.text{fill:none;stroke-width:0;}#mermaid-svg-MTAsSupYDSaVkI7O .icon-shape,#mermaid-svg-MTAsSupYDSaVkI7O .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-MTAsSupYDSaVkI7O .icon-shape p,#mermaid-svg-MTAsSupYDSaVkI7O .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-MTAsSupYDSaVkI7O .icon-shape .label rect,#mermaid-svg-MTAsSupYDSaVkI7O .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-MTAsSupYDSaVkI7O .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-MTAsSupYDSaVkI7O .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-MTAsSupYDSaVkI7O :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} B 的内存

(B 已死亡)
空闲内存池
E 接管这块内存

Inplace vs Sharing 的区别:

Inplace Sharing
关系 输入和输出是同一个算子 两个不相关的张量(生命周期不重叠)
前提 算子支持、原值不再需要 前一个已死、大小匹配
典型 ReLU、add_、mul_ 层与层之间的缓冲区复用

两者本质是同一件事 :在时间上不重叠的活人,可以睡同一张床。Inplace 是"极短间隔"的特例,Sharing 是"跨算子"的一般情况。


四、标签 + 内存池:一个可执行的分配算法

原文用一张分步图讲了一个算法,核心是 "初始化时统计使用次数 + 用 tag 标记归属"。

4.1 算法流程

#mermaid-svg-rTKfBiK7ymLdg3jE{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-rTKfBiK7ymLdg3jE .error-icon{fill:#552222;}#mermaid-svg-rTKfBiK7ymLdg3jE .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-rTKfBiK7ymLdg3jE .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-rTKfBiK7ymLdg3jE .marker{fill:#333333;stroke:#333333;}#mermaid-svg-rTKfBiK7ymLdg3jE .marker.cross{stroke:#333333;}#mermaid-svg-rTKfBiK7ymLdg3jE svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-rTKfBiK7ymLdg3jE p{margin:0;}#mermaid-svg-rTKfBiK7ymLdg3jE .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster-label text{fill:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster-label span{color:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster-label span p{background-color:transparent;}#mermaid-svg-rTKfBiK7ymLdg3jE .label text,#mermaid-svg-rTKfBiK7ymLdg3jE span{fill:#333;color:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE .node rect,#mermaid-svg-rTKfBiK7ymLdg3jE .node circle,#mermaid-svg-rTKfBiK7ymLdg3jE .node ellipse,#mermaid-svg-rTKfBiK7ymLdg3jE .node polygon,#mermaid-svg-rTKfBiK7ymLdg3jE .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-rTKfBiK7ymLdg3jE .rough-node .label text,#mermaid-svg-rTKfBiK7ymLdg3jE .node .label text,#mermaid-svg-rTKfBiK7ymLdg3jE .image-shape .label,#mermaid-svg-rTKfBiK7ymLdg3jE .icon-shape .label{text-anchor:middle;}#mermaid-svg-rTKfBiK7ymLdg3jE .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-rTKfBiK7ymLdg3jE .rough-node .label,#mermaid-svg-rTKfBiK7ymLdg3jE .node .label,#mermaid-svg-rTKfBiK7ymLdg3jE .image-shape .label,#mermaid-svg-rTKfBiK7ymLdg3jE .icon-shape .label{text-align:center;}#mermaid-svg-rTKfBiK7ymLdg3jE .node.clickable{cursor:pointer;}#mermaid-svg-rTKfBiK7ymLdg3jE .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-rTKfBiK7ymLdg3jE .arrowheadPath{fill:#333333;}#mermaid-svg-rTKfBiK7ymLdg3jE .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-rTKfBiK7ymLdg3jE .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-rTKfBiK7ymLdg3jE .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-rTKfBiK7ymLdg3jE .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-rTKfBiK7ymLdg3jE .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-rTKfBiK7ymLdg3jE .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster text{fill:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE .cluster span{color:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-rTKfBiK7ymLdg3jE .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-rTKfBiK7ymLdg3jE rect.text{fill:none;stroke-width:0;}#mermaid-svg-rTKfBiK7ymLdg3jE .icon-shape,#mermaid-svg-rTKfBiK7ymLdg3jE .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-rTKfBiK7ymLdg3jE .icon-shape p,#mermaid-svg-rTKfBiK7ymLdg3jE .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-rTKfBiK7ymLdg3jE .icon-shape .label rect,#mermaid-svg-rTKfBiK7ymLdg3jE .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-rTKfBiK7ymLdg3jE .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-rTKfBiK7ymLdg3jE .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-rTKfBiK7ymLdg3jE :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 初始化:遍历计算图

统计每个张量的使用次数
Step 1:为 B 分配 tag
Step 2:为 C 分配新 tag

B 还活着 → 不能 Inplace
Step 3:计算 F

F 用完 B → B 生命周期结束

B 的内存进入空闲队列
Step 4:C 结束,内存入池

执行 E → 从池中取 B 的块

给 E 做 Memory Sharing

同时更新 tag 标记
Step 5:E 之后的操作

可复用 E 的内存 → Inplace

用时间线看更清楚:
#mermaid-svg-K2a3IbREyEDnovNu{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-K2a3IbREyEDnovNu .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-K2a3IbREyEDnovNu .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-K2a3IbREyEDnovNu .error-icon{fill:#552222;}#mermaid-svg-K2a3IbREyEDnovNu .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-K2a3IbREyEDnovNu .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-K2a3IbREyEDnovNu .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-K2a3IbREyEDnovNu .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-K2a3IbREyEDnovNu .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-K2a3IbREyEDnovNu .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-K2a3IbREyEDnovNu .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-K2a3IbREyEDnovNu .marker{fill:#333333;stroke:#333333;}#mermaid-svg-K2a3IbREyEDnovNu .marker.cross{stroke:#333333;}#mermaid-svg-K2a3IbREyEDnovNu svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-K2a3IbREyEDnovNu p{margin:0;}#mermaid-svg-K2a3IbREyEDnovNu .mermaid-main-font{font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-K2a3IbREyEDnovNu .exclude-range{fill:#eeeeee;}#mermaid-svg-K2a3IbREyEDnovNu .section{stroke:none;opacity:0.2;}#mermaid-svg-K2a3IbREyEDnovNu .section0{fill:rgba(102, 102, 255, 0.49);}#mermaid-svg-K2a3IbREyEDnovNu .section2{fill:#fff400;}#mermaid-svg-K2a3IbREyEDnovNu .section1,#mermaid-svg-K2a3IbREyEDnovNu .section3{fill:white;opacity:0.2;}#mermaid-svg-K2a3IbREyEDnovNu .sectionTitle0{fill:#333;}#mermaid-svg-K2a3IbREyEDnovNu .sectionTitle1{fill:#333;}#mermaid-svg-K2a3IbREyEDnovNu .sectionTitle2{fill:#333;}#mermaid-svg-K2a3IbREyEDnovNu .sectionTitle3{fill:#333;}#mermaid-svg-K2a3IbREyEDnovNu .sectionTitle{text-anchor:start;font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-K2a3IbREyEDnovNu .grid .tick{stroke:lightgrey;opacity:0.8;shape-rendering:crispEdges;}#mermaid-svg-K2a3IbREyEDnovNu .grid .tick text{font-family:"trebuchet ms",verdana,arial,sans-serif;fill:#333;}#mermaid-svg-K2a3IbREyEDnovNu .grid path{stroke-width:0;}#mermaid-svg-K2a3IbREyEDnovNu .today{fill:none;stroke:red;stroke-width:2px;}#mermaid-svg-K2a3IbREyEDnovNu .task{stroke-width:2;}#mermaid-svg-K2a3IbREyEDnovNu .taskText{text-anchor:middle;font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutsideRight{fill:black;text-anchor:start;font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutsideLeft{fill:black;text-anchor:end;}#mermaid-svg-K2a3IbREyEDnovNu .task.clickable{cursor:pointer;}#mermaid-svg-K2a3IbREyEDnovNu .taskText.clickable{cursor:pointer;fill:#003163!important;font-weight:bold;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutsideLeft.clickable{cursor:pointer;fill:#003163!important;font-weight:bold;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutsideRight.clickable{cursor:pointer;fill:#003163!important;font-weight:bold;}#mermaid-svg-K2a3IbREyEDnovNu .taskText0,#mermaid-svg-K2a3IbREyEDnovNu .taskText1,#mermaid-svg-K2a3IbREyEDnovNu .taskText2,#mermaid-svg-K2a3IbREyEDnovNu .taskText3{fill:white;}#mermaid-svg-K2a3IbREyEDnovNu .task0,#mermaid-svg-K2a3IbREyEDnovNu .task1,#mermaid-svg-K2a3IbREyEDnovNu .task2,#mermaid-svg-K2a3IbREyEDnovNu .task3{fill:#8a90dd;stroke:#534fbc;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutside0,#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutside2{fill:black;}#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutside1,#mermaid-svg-K2a3IbREyEDnovNu .taskTextOutside3{fill:black;}#mermaid-svg-K2a3IbREyEDnovNu .active0,#mermaid-svg-K2a3IbREyEDnovNu .active1,#mermaid-svg-K2a3IbREyEDnovNu .active2,#mermaid-svg-K2a3IbREyEDnovNu .active3{fill:#bfc7ff;stroke:#534fbc;}#mermaid-svg-K2a3IbREyEDnovNu .activeText0,#mermaid-svg-K2a3IbREyEDnovNu .activeText1,#mermaid-svg-K2a3IbREyEDnovNu .activeText2,#mermaid-svg-K2a3IbREyEDnovNu .activeText3{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .done0,#mermaid-svg-K2a3IbREyEDnovNu .done1,#mermaid-svg-K2a3IbREyEDnovNu .done2,#mermaid-svg-K2a3IbREyEDnovNu .done3{stroke:grey;fill:lightgrey;stroke-width:2;}#mermaid-svg-K2a3IbREyEDnovNu .doneText0,#mermaid-svg-K2a3IbREyEDnovNu .doneText1,#mermaid-svg-K2a3IbREyEDnovNu .doneText2,#mermaid-svg-K2a3IbREyEDnovNu .doneText3{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .doneText0.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneText0.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneText1.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneText1.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneText2.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneText2.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneText3.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneText3.taskTextOutsideRight{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .crit0,#mermaid-svg-K2a3IbREyEDnovNu .crit1,#mermaid-svg-K2a3IbREyEDnovNu .crit2,#mermaid-svg-K2a3IbREyEDnovNu .crit3{stroke:#ff8888;fill:red;stroke-width:2;}#mermaid-svg-K2a3IbREyEDnovNu .activeCrit0,#mermaid-svg-K2a3IbREyEDnovNu .activeCrit1,#mermaid-svg-K2a3IbREyEDnovNu .activeCrit2,#mermaid-svg-K2a3IbREyEDnovNu .activeCrit3{stroke:#ff8888;fill:#bfc7ff;stroke-width:2;}#mermaid-svg-K2a3IbREyEDnovNu .doneCrit0,#mermaid-svg-K2a3IbREyEDnovNu .doneCrit1,#mermaid-svg-K2a3IbREyEDnovNu .doneCrit2,#mermaid-svg-K2a3IbREyEDnovNu .doneCrit3{stroke:#ff8888;fill:lightgrey;stroke-width:2;cursor:pointer;shape-rendering:crispEdges;}#mermaid-svg-K2a3IbREyEDnovNu .milestone{transform:rotate(45deg) scale(0.8,0.8);}#mermaid-svg-K2a3IbREyEDnovNu .milestoneText{font-style:italic;}#mermaid-svg-K2a3IbREyEDnovNu .doneCritText0,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText1,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText2,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText3{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .doneCritText0.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText0.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText1.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText1.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText2.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText2.taskTextOutsideRight,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText3.taskTextOutsideLeft,#mermaid-svg-K2a3IbREyEDnovNu .doneCritText3.taskTextOutsideRight{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .vert{stroke:navy;}#mermaid-svg-K2a3IbREyEDnovNu .vertText{font-size:15px;text-anchor:middle;fill:navy!important;}#mermaid-svg-K2a3IbREyEDnovNu .activeCritText0,#mermaid-svg-K2a3IbREyEDnovNu .activeCritText1,#mermaid-svg-K2a3IbREyEDnovNu .activeCritText2,#mermaid-svg-K2a3IbREyEDnovNu .activeCritText3{fill:black!important;}#mermaid-svg-K2a3IbREyEDnovNu .titleText{text-anchor:middle;font-size:18px;fill:#333;font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-K2a3IbREyEDnovNu :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 0 0 1 1 2 2 3 3 4 4 5 5 6 B 使用中 C 使用中 F 使用中 E 复用(Sharing) 缓冲区1 缓冲区2 缓冲区3 张量生命周期与缓冲区复用

注意原文的提醒 :这是**预执行(编译期模拟)**得到的分配方案,不是真实运行时。真实分配还要考虑对齐、内存池、碎片和运行时动态形状。

4.2 为什么这个算法像寄存器分配

原文有一句点睛:

AI 编译器的内存分配算法与传统编译器的寄存器分配非常相似,可以借鉴许多思想。

传统编译 AI 编译器
虚拟寄存器 张量(Tensor)
物理寄存器 显存块(Buffer)
寄存器溢出(spill) 换出到 CPU 内存 / 重计算
活跃变量分析(liveness) 张量生命周期分析
干涉图(interference graph) 生命周期重叠图
图着色分配 缓冲区分配

形式化地说:

  1. 每个张量有一个生命周期区间 t def , t last-use t_{\\text{def}}, t_{\\text{last-use}} tdef,tlast-use
  2. 两个张量区间重叠 → 不能共用同一块内存
  3. 构造干涉图,求最少颜色数(= 最少缓冲区数)
  4. 峰值内存 = 任一时刻所有存活张量大小之和的最大值

M peak = max ⁡ t ∑ v ∈ live ( t ) size ( v ) M_{\text{peak}} = \max_{t} \sum_{v \in \text{live}(t)} \text{size}(v) Mpeak=tmaxv∈live(t)∑size(v)

大白话 :同一时刻活着的人,不能挤一张床。目标是让"同时活着的人"尽量少、尽量小。

4.3 内存池:别每次都 malloc

真实编译器通常不会对每个张量调用一次 malloc:

做法 好处
预分配大块 + 池化管理 避免频繁系统调用和碎片
按 size class 分桶 复用率高、查找快
对齐到固定粒度(如 256 B / 512 B) 满足硬件访存对齐要求(呼应布局那篇)
及时归还到池 让后续张量能立刻复用

五、并行分配:别让"省内存"把并行性省没了

原文最后讲了一个非常容易被忽略的问题:内存复用可能引入假的依赖,从而限制并行。

5.1 两种方案,同样的串行语义

同一段程序 A[1] → A[8],两种内存分配方案:
#mermaid-svg-koDJ1Ne01LXVlzRu{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-koDJ1Ne01LXVlzRu .error-icon{fill:#552222;}#mermaid-svg-koDJ1Ne01LXVlzRu .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-koDJ1Ne01LXVlzRu .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-koDJ1Ne01LXVlzRu .marker{fill:#333333;stroke:#333333;}#mermaid-svg-koDJ1Ne01LXVlzRu .marker.cross{stroke:#333333;}#mermaid-svg-koDJ1Ne01LXVlzRu svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-koDJ1Ne01LXVlzRu p{margin:0;}#mermaid-svg-koDJ1Ne01LXVlzRu .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster-label text{fill:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster-label span{color:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster-label span p{background-color:transparent;}#mermaid-svg-koDJ1Ne01LXVlzRu .label text,#mermaid-svg-koDJ1Ne01LXVlzRu span{fill:#333;color:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu .node rect,#mermaid-svg-koDJ1Ne01LXVlzRu .node circle,#mermaid-svg-koDJ1Ne01LXVlzRu .node ellipse,#mermaid-svg-koDJ1Ne01LXVlzRu .node polygon,#mermaid-svg-koDJ1Ne01LXVlzRu .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-koDJ1Ne01LXVlzRu .rough-node .label text,#mermaid-svg-koDJ1Ne01LXVlzRu .node .label text,#mermaid-svg-koDJ1Ne01LXVlzRu .image-shape .label,#mermaid-svg-koDJ1Ne01LXVlzRu .icon-shape .label{text-anchor:middle;}#mermaid-svg-koDJ1Ne01LXVlzRu .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-koDJ1Ne01LXVlzRu .rough-node .label,#mermaid-svg-koDJ1Ne01LXVlzRu .node .label,#mermaid-svg-koDJ1Ne01LXVlzRu .image-shape .label,#mermaid-svg-koDJ1Ne01LXVlzRu .icon-shape .label{text-align:center;}#mermaid-svg-koDJ1Ne01LXVlzRu .node.clickable{cursor:pointer;}#mermaid-svg-koDJ1Ne01LXVlzRu .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-koDJ1Ne01LXVlzRu .arrowheadPath{fill:#333333;}#mermaid-svg-koDJ1Ne01LXVlzRu .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-koDJ1Ne01LXVlzRu .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-koDJ1Ne01LXVlzRu .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-koDJ1Ne01LXVlzRu .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-koDJ1Ne01LXVlzRu .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-koDJ1Ne01LXVlzRu .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster text{fill:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu .cluster span{color:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-koDJ1Ne01LXVlzRu .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-koDJ1Ne01LXVlzRu rect.text{fill:none;stroke-width:0;}#mermaid-svg-koDJ1Ne01LXVlzRu .icon-shape,#mermaid-svg-koDJ1Ne01LXVlzRu .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-koDJ1Ne01LXVlzRu .icon-shape p,#mermaid-svg-koDJ1Ne01LXVlzRu .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-koDJ1Ne01LXVlzRu .icon-shape .label rect,#mermaid-svg-koDJ1Ne01LXVlzRu .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-koDJ1Ne01LXVlzRu .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-koDJ1Ne01LXVlzRu .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-koDJ1Ne01LXVlzRu :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 方案乙:内存稍多,但可并行
A1
A2
A5
A3
A6
方案甲:省内存,但依赖多
A1
A2
A3
A4
A5
A6
A7
A8

如果按单线程顺序执行,两种方案结果都对。但:

方案 后果
甲 A5 复用了 A2 的缓冲区 → 必须等 A2 的所有消费者读完 → 引入 WAW/RAW 假依赖 → A2 和 A5 不能并行
乙 A5 用独立缓冲区 → A2 和 A5 可以并行 → 吞吐更高

#mermaid-svg-pRikElEgkFe3k9na{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-pRikElEgkFe3k9na .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-pRikElEgkFe3k9na .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-pRikElEgkFe3k9na .error-icon{fill:#552222;}#mermaid-svg-pRikElEgkFe3k9na .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-pRikElEgkFe3k9na .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-pRikElEgkFe3k9na .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-pRikElEgkFe3k9na .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-pRikElEgkFe3k9na .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-pRikElEgkFe3k9na .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-pRikElEgkFe3k9na .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-pRikElEgkFe3k9na .marker{fill:#333333;stroke:#333333;}#mermaid-svg-pRikElEgkFe3k9na .marker.cross{stroke:#333333;}#mermaid-svg-pRikElEgkFe3k9na svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-pRikElEgkFe3k9na p{margin:0;}#mermaid-svg-pRikElEgkFe3k9na .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-pRikElEgkFe3k9na .cluster-label text{fill:#333;}#mermaid-svg-pRikElEgkFe3k9na .cluster-label span{color:#333;}#mermaid-svg-pRikElEgkFe3k9na .cluster-label span p{background-color:transparent;}#mermaid-svg-pRikElEgkFe3k9na .label text,#mermaid-svg-pRikElEgkFe3k9na span{fill:#333;color:#333;}#mermaid-svg-pRikElEgkFe3k9na .node rect,#mermaid-svg-pRikElEgkFe3k9na .node circle,#mermaid-svg-pRikElEgkFe3k9na .node ellipse,#mermaid-svg-pRikElEgkFe3k9na .node polygon,#mermaid-svg-pRikElEgkFe3k9na .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-pRikElEgkFe3k9na .rough-node .label text,#mermaid-svg-pRikElEgkFe3k9na .node .label text,#mermaid-svg-pRikElEgkFe3k9na .image-shape .label,#mermaid-svg-pRikElEgkFe3k9na .icon-shape .label{text-anchor:middle;}#mermaid-svg-pRikElEgkFe3k9na .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-pRikElEgkFe3k9na .rough-node .label,#mermaid-svg-pRikElEgkFe3k9na .node .label,#mermaid-svg-pRikElEgkFe3k9na .image-shape .label,#mermaid-svg-pRikElEgkFe3k9na .icon-shape .label{text-align:center;}#mermaid-svg-pRikElEgkFe3k9na .node.clickable{cursor:pointer;}#mermaid-svg-pRikElEgkFe3k9na .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-pRikElEgkFe3k9na .arrowheadPath{fill:#333333;}#mermaid-svg-pRikElEgkFe3k9na .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-pRikElEgkFe3k9na .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-pRikElEgkFe3k9na .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-pRikElEgkFe3k9na .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-pRikElEgkFe3k9na .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-pRikElEgkFe3k9na .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-pRikElEgkFe3k9na .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-pRikElEgkFe3k9na .cluster text{fill:#333;}#mermaid-svg-pRikElEgkFe3k9na .cluster span{color:#333;}#mermaid-svg-pRikElEgkFe3k9na div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-pRikElEgkFe3k9na .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-pRikElEgkFe3k9na rect.text{fill:none;stroke-width:0;}#mermaid-svg-pRikElEgkFe3k9na .icon-shape,#mermaid-svg-pRikElEgkFe3k9na .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-pRikElEgkFe3k9na .icon-shape p,#mermaid-svg-pRikElEgkFe3k9na .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-pRikElEgkFe3k9na .icon-shape .label rect,#mermaid-svg-pRikElEgkFe3k9na .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-pRikElEgkFe3k9na .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-pRikElEgkFe3k9na .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-pRikElEgkFe3k9na :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 是
否
想并行执行 A2 与 A5
两者共用缓冲区?
必须串行

(先读完再覆盖)
可以并行

核心权衡 :内存复用是在"时间维度"上省钱,但可能牺牲"空间维度"上的并行。

在 GPU/NPU 上,这个代价往往不划算 ------因为并行度直接决定吞吐。所以成熟的编译器会在分配时识别并行分支,给它们分配不同缓冲区,哪怕多花一点内存。

5.2 实际怎么权衡

场景 倾向
单流、显存紧张(如大模型推理) 激进复用,优先省内存
多流 / 多分支并行 保守复用,优先保并行
训练(激活值占大头) 复用 + 重计算组合
小算子密集(element-wise 链) Inplace 收益最大

六、MobileNetV2 的例子:优化前后差多少

原文给了一张对比图(上半:未优化;下半:优化后):

指标 优化前 优化后
内存占用 高 明显下降
碎片数量 多(大量小块) 显著减少
不同算子的内存需求 各自申请 整合复用

为什么会这样? MobileNetV2 大量使用 深度可分离卷积 + 逐元素激活:

  • 层数多、每层张量不大 → 碎片化严重
  • 大量 element-wise 算子 → Inplace 机会多
  • 结构规整、生命周期清晰 → 复用率高

这正是编译器擅长的场景:人手工算不过来,但算法可以全局看。


七、小结

# 结论
1 内存分配决策在前端、执行在后端;前端负责生命周期分析与复用方案
2 算力涨得比显存快,显存是稀缺资源;40 GB 训练显存里权重只占很小一部分
3 按生命周期分更准确:持久 (权重/梯度/优化器状态/常量)vs 瞬态(激活/workspace/临时)
4 原文把"输出"算作静态内存是错的(作者自己也标注了);输出应由调用方持有
5 梯度检查点把"存激活 "换成"重算激活",是训练省显存的主力手段
6 Adam 混合精度训练约需 16P 字节(权重 2P + 梯度 2P + master 4P + m 4P + v 4P)
7 Inplace:element-wise + 原值不再需要 + 反向不需要原值,才能原地覆盖
8 Sharing:前一个张量已死 + 大小匹配 → 后一个直接接管其缓冲区
9 复用算法本质是寄存器分配:生命周期区间 → 干涉图 → 最少缓冲区
10 峰值内存 = 同时存活张量大小之和的最大值
11 内存复用可能引入假依赖、杀掉并行;并行分支应分配独立缓冲区
12 真实实现要配合内存池 + 对齐 + 分桶,避免频繁 malloc 和碎片

附:对原文的修正

# 原文 修正
1 静态内存 = 权重 + 值节点 + 输出 输出不属于静态内存(原文自己标注"肯定不对");且与"输出张量"重复分类
2 未提优化器状态/动量 训练显存大头之一:Adam 的 m、v 各 4P,混合精度还要 master weights
3 "值节点(Value Node)" 通常指图中常量节点;可折叠的会被常量折叠消除,不可折叠的才常驻
4 GPU 表 CUDA 核心行 7 个值对 6 列 已修正为 6 个值;A100 的 INT8 不是 N/A(312 TOPS 稠密)
5 "量化训练 Quantification" 应为 Quantization
6 "以空间换内存" 表述含糊,本质是 offload:用 CPU 内存/磁盘换显存
7 静态/动态划分的边界 建议改用**生命周期(持久/瞬态)**作为判据,更不容易出错
8 算法示例未说明是编译期模拟 原文有提"预执行",本文补强:这不是运行时真实分配

参考资料

相关推荐
又蓝1 小时前
【元宝搜索开发者征文大赛】用 FastAPI + Vue3 做了个 AI 炒股工具,接了六个分析师和腾讯元宝搜索
人工智能·python·ai·fastapi
LittroInno1 小时前
边缘AI模组选型:把视频带宽、压缩和推理吞吐分别算清
人工智能·音视频·边缘计算·视频编码·模组选型·视频带宽·tofu6
niucloud-admin1 小时前
NIUCLOUD PHP SAAS上门家政使用文档:机构报表
前端·人工智能
AI小百科1 小时前
Embedding 是什么?为什么它是 AI 的“翻译官“
人工智能·embedding
进击的横打1 小时前
【人工智能】AI 落地:先理流程,再做 Agent
人工智能
虎鲸不是鱼1 小时前
AI 自动化逆向(二十二)换行比峰值法:让 AI 自己发现图像的宽度
人工智能·自动化·大模型·逆向·ai逆向
qingyunliushuiyu2 小时前
企业选数据BI AI服务商时,需关注哪些核心能力?
人工智能·经验分享
2601_970703832 小时前
中国除湿机与湿度控制行业研究报告
java·大数据·人工智能
天天代码码天天2 小时前
把自己的字,变成电脑里的字体:我做了一个「手写字体工坊」
人工智能