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) | 生命周期重叠图 |
| 图着色分配 | 缓冲区分配 |
形式化地说:
- 每个张量有一个生命周期区间 t def , t last-use t_{\\text{def}}, t_{\\text{last-use}} tdef,tlast-use
- 两个张量区间重叠 → 不能共用同一块内存
- 构造干涉图,求最少颜色数(= 最少缓冲区数)
- 峰值内存 = 任一时刻所有存活张量大小之和的最大值
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 | 算法示例未说明是编译期模拟 | 原文有提"预执行",本文补强:这不是运行时真实分配 |
参考资料
- 原文:AISystem / 03Compiler / 03Frontend / 06Memory.md(作者 ZOMI)
- 配套视频:BV1nM411879s
- 寄存器分配与图着色:Engineering a Compiler 第 11 章
- 梯度检查点:https://pytorch.org/docs/stable/checkpoint.html
- ZeRO(切分优化器状态/梯度/参数):https://arxiv.org/abs/1910.02054
- 混合精度训练:https://arxiv.org/abs/1710.03740