1. 动态计算图PyTorch 为什么偏要在运行时现搭 DAG不少同学从 TensorFlow 1.x 时代转过来用 PyTorch 时最不适应的就是怎么没看到图结构在 TensorFlow 1.x 里你得先用tf.Graph定义好一张静态计算图再丢给Session去跑而 PyTorch 上手就是直接写张量运算看起来跟写普通 Python 脚本几乎没区别。直到你第一次打印一个张量的grad_fn属性才发现奥妙藏在背后。1.1 计算图不是画出来的是记下来的静态计算图的构建方式是先搭图后执行图结构一旦定义基本固定改动特别麻烦。PyTorch 走的完全是另一条路它在每次前向传播执行张量运算时动态记录操作历史自动把算子之间依赖关系编织成一张有向无环图DAG。这张图不是预先画好存在某个对象里的而是通过一个个 Tensor 的grad_fn属性层层串联起来的。举个例子我用最朴素的三行代码演示这一点import torch x torch.tensor([2.0], requires_gradTrue) z x ** 2 # z x^2 y 3 * z 5 # y 3x^2 5 print(x.grad_fn) # None因为 x 是叶子节点 print(z.grad_fn) # PowBackward0 object at 0x... print(y.grad_fn) # AddBackward0 object at 0x...z.grad_fn指向的是生成 z 的那个操作的逆向函数入口。这个对象内部还保留着指向输入张量x的引用以及操作所需的元数据。PyTorch 就是靠这种链式结构把整个前向传播过程组织成一张可以从输出一路回溯到输入的图。我把动态图和静态图做了个直观对比对比维度动态图PyTorch静态图TF1.x图构建时机每次前向传播运行时动态构建先定义后执行调试体验可以用 print、断点、if/else 原生 Python 控制流需要把控制流改成图节点调试困难性能优化空间编译优化难但灵活性极高可做图优化、算子融合典型适用场景研究、动态结构模型、RL生产部署、静态推理为什么 PyTorch 选择了磁带式自动微分这条路线核心原因是研究和实验中模型结构经常需要调整网络层数、分支结构甚至损失函数都可能前一天才改过。动态图意味着你完全可以用if语句判断输入长度来决定走哪个分支因为图就是跟着 Python 执行轨迹现搭的不需要为难你。1.2 磁带式记录DAG 的建立时机和内部关联PyTorch 的自动微分原理其实可以类比成录音机或者磁带你在前向传播过程中做的每一步张量运算都像在磁带上录下操作类型、输入张量、输出张量这三要素。反向传播时自动微分系统把磁带倒过来放沿着记录倒序执行每个操作对应的求导规则把梯度逐级传回去。再来看看 DAG 在内存里到底是什么。执行一段前向传播代码时PyTorch 会在每个需要梯度的张量之间建立边关系。拿最常见的线性层举例x torch.randn(64, 512, requires_gradTrue) w torch.randn(512, 256, requires_gradTrue) b torch.randn(256, requires_gradTrue) h x w # 矩阵乘 h h b # 广播加法 out h.tanh() # 激活函数这四行代码运行时内存中形成的 DAG 大致是这样的结构x(requires_gradTrue) w(requires_gradTrue) b(requires_gradTrue) \ / / \ / / \ / / \ / / \ / / MatMulBackward / \ / \ / \ / AddBackward / \ / \ / \ / TanhBackward | out每个反向函数对象MatMulBackward、AddBackward、TanhBackward内部都保存了前向计算时的输入。这其实是一个非常重要的细节DAG 不仅保存了结构还保存了前向运算的中间结果。正因如此反向传播时才能直接复用前向的值来算梯度比如tanh激活函数的梯度需要用到前向的out值本身。这些被保存下来的中间张量正是 Activation 显存开销的根源。1.3 叶子节点与非叶子节点的命运差异凡是requires_gradTrue且由用户直接创建、不依赖其他张量的 Tensor被称为叶子节点Leaf Tensor。非叶子节点则是由运算生成的中间张量。它们在反向传播结束后的梯度处理上待遇完全不同叶子节点反向传播完成后梯度会累积在.grad属性里留给优化器更新参数使用。如果你连续调用两次loss.backward()叶子节点的梯度是叠加的所以训练循环里每步都得手动optimizer.zero_grad()。非叶子节点反向传播完梯度只作为中间传递的桥梁用完后如果不做特殊保存会直接被丢弃释放。你要是试图打印中间张量的.grad得到的往往是None。非叶子节点如果想保留梯度需要调用.retain_grad()方法提前声明。这个机制我刚开始学 PyTorch 时经常忽略直到有一天我想可视化某层特征图的梯度打印出来全是None才意识到 PyTorch 为了省内存放弃了自动保存非叶子节点梯度的能力。2. Autograd 反向求导机制链式法则如何在 DAG 上传播反向传播表面看只是调用一下.backward()但这一行调用背后隐藏着一套完整的任务调度系统从输出节点出发按照 DAG 的拓扑逆序遍历所有反向节点依次执行局部链式求导把梯度向叶子节点方向逐级传递。2.1 从链式法则到梯度传播路径复杂模型的链式法则本质上是分段求导逐级相乘。PyTorch 对每个算子都预先注册了局部求导规则比如乘法的反向规则、卷积的反向规则。反向传播时Autograd 引擎只需要把上游传来的梯度与当前算子的局部雅可比矩阵相乘就得到这个算子各输入的梯度贡献。拿最基础的复合函数y f(g(x))举例dy/dx dy/du * du/dx其中u g(x)。在 PyTorch 中这一过程被拆成两次反向节点的传播x torch.tensor([3.0], requires_gradTrue) u 2 * x 1 # u g(x)grad_fnAddBackward y u ** 2 # y f(u)grad_fnPowBackward y.backward() print(x.grad) # 12.0即 dy/du6, du/dx2, 乘积 12y.backward()执行时Autograd 引擎首先计算dy/du也就是 2u 6然后把梯度 6 传给u的反向节点AddBackward再计算du/dx 2二者相乘得到 12写入x.grad。这就是链式法则在 DAG 上传播的完整路径。2.2 Autograd 引擎的任务调度策略PyTorch 的 Autograd 引擎在反向传播时采用了基于线程池的异步执行策略。它会从输出节点开始把无依赖的梯度计算任务放入队列多个工作线程并行处理同一层级的梯度任务当一个节点的所有输入梯度都计算完成之后才能触发该节点的梯度计算。这样一层一层解锁下去直到所有叶子节点的梯度都计算完毕。这也是为什么在反向传播过程中我们有时会看到 GPU 利用率上的锯齿波动——不同层因为计算量差异在各个阶段的表现节奏天然不一样而 Autograd 引擎是动态调度的哪里的依赖先满足就先算哪里。老版本 PyTorch 在复杂模型上有时会因为线程调度的开销卡顿1.9 之后引入了一种基于链表结构的新调度实现显著降低了反向传播的线程同步开销。实测在包含大量小算子的模型上比如密集的 attention 结构新老版本的 backward 时间能差出 10%~20%。2.3 retain_graph、detach 与梯度累积的关系调用.backward()后Autograd 引擎会把整张计算图的中间缓存全部释放。如果对同一个 loss 连续调用两次.backward()比如某些 GAN 训练代码里 generator 和 discriminator 分开更新会直接报错RuntimeError: Trying to backward through the graph a second time。解决方式有两个方向在backward()里传入retain_graphTrue强制保留计算图缓存但这样内存开销和第一次前向持平。如果你的计算图很大这一招会让显存雪上加霜。更推荐的做法是搞明白梯度累积的语义如果只是想把多批数据的梯度加到一起完全可以用.detach()切断分支或者分多次调用 backward 而每次使用不同的子图输出。比如 GAN 里 discriminator 的 loss 和 generator 的 loss 通常来自不同的计算路径只需要保证共享部分参数的计算图在当前 batch 内只 backward 一次即可。还有一个高频踩坑点把 loss 的 item 标量值拿来计算后又参与梯度计算导致图被意外保留。正确做法就是把数值统计的部分.detach()出来不要让无关的 Python 运算串进计算图。3. 显存生命周期前向、反向、优化器更新显存究竟装了什么很多人用 PyTorch 训练时只看一个整体显存占用数字其实训练一步过程中显存是被不同角色瓜分的模型参数、优化器动量、前向中间激活、反向梯度、CUDA context 和计算临时缓存。搞清楚每个时刻显存里都有谁是显存优化的重要基础。3.1 一个典型训练步骤中的显存阶段变化以最常见的一个 step 为例假设你用了一张 12GB 的显卡训练 ResNet-50输入 batch size 为 32。我用torch.cuda.memory_allocated()在两个关键时点打点观测看到的现象基本如下model torchvision.models.resnet50().cuda() optimizer torch.optim.SGD(model.parameters(), lr0.01) dummy_input torch.randn(32, 3, 224, 224).cuda() torch.cuda.synchronize() print(初始显存占用:, torch.cuda.memory_allocated() / 1024**2, MB) output model(dummy_input) loss output.sum() torch.cuda.synchronize() print(前向完成显存:, torch.cuda.memory_allocated() / 1024**2, MB) loss.backward() torch.cuda.synchronize() print(反向完成显存:, torch.cuda.memory_allocated() / 1024**2, MB) optimizer.step() torch.cuda.synchronize() print(优化器更新后:, torch.cuda.memory_allocated() / 1024**2, MB)实测规律非常稳定初始显存只包含模型参数和 CUDA context大概 100MB 到 400MB 的数量级。前向完成显存猛增到几个 GB多出来的全部是各层输出的 Activation 中间激活值。反向完成相比前向完成时没有明显变化甚至可能小幅度下降因为反向计算过程中读取并释放了部分 Activation 缓存但同时产生了各参数的梯度。优化器更新后显存基本保持不变参数梯度不会因为optimizer.step()而释放除非手动清空因为梯度要保留到下一个zero_grad()。这一个周期就把显存生命周期分得明明白白参数和梯度占的只是冰山一角前向传播时堆积起来的 Activation 才是显存水位暴涨的主导因素。3.2 CUDA 缓存分配器的假释放现象还有一个让很多人迷惑的现象你用nvidia-smi看显存占用明明很高但代码里torch.cuda.memory_allocated()显示占用很低。这是因为 PyTorch 为了减少频繁调用 CUDA 底层显存分配接口的系统开销默认使用缓存分配器。Tensor 被 Python 端释放后显存块并不会真的还给显存驱动而是留在 PyTorch 的缓存池里等待复用。这个设计的代价是你在nvidia-smi中看到的显存剩余其实不准确。有一次我跑一个需要动态 batch size 的推理服务每个请求的张量形状都在变化导致缓存池里碎片化严重发生了即使memory_allocated不高但依然 OOM 的情况最后只能通过PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True缓解。顺带提一个排查技巧出现 OOM 时先运行torch.cuda.memory_summary()它会给出当前显存池里缓存块的大小分布和分配器行为明细远比nvidia-smi的信息有诊断价值。3.3 梯度传播结束后的释放时机和内存复用反向传播结束后非叶子节点的 Activation 缓存会立刻被释放回缓存池这些显存块可以马上被下一轮前向传播复用这就是为什么连续多步训练时显存占用基本持平而不会持续上涨。反过来说如果某个中间结果因为被外部变量引用而迟迟得不到释放显存占用就会异常增高。这里有三个高频错误optimizer.zero_grad()放在 backward 前面而不是每步更新后调用导致梯度叠加。在循环里把每个 batch 的输出 append 到一个 Python list计算图一直没释放list 越界增长显存逐步耗尽。开了torch.no_grad()还挂requires_gradTrue张量做验证循环虽然不会建图但历史张量被保留也会累积。4. Activation 显存优化从 checkpointing 到混合精度的完整手段既然 Activation 是显存水位的绝对主导者那如何压低它就成了训练大规模模型的必修课。业内最核心的思路说白了只有两个方向减少同时保留在显存里的中间激活或者降低每个激活的存储精度。4.1 Activation Checkpointing用计算换显存Activation Checkpointing也叫梯度检查点的原理非常直白在反向传播时本来需要前向保存的每个中间结果现在只保留若干个检查点位置的输出其余中间结果是反向传播时再临时重新前向计算的。PyTorch 提供了一个现成的 APIfrom torch.utils.checkpoint import checkpoint def transformer_block(x): # 一些计算密集型结构 return x x checkpoint(transformer_block, x)这里需要注意checkpoint 不是免费午餐。它额外触发一次前向计算所以训练时间通常会增加 20%~30%。但对于那些显存已经告急的大模型这个时间换空间的买卖非常划算可以让原本完全跑不动的模型训练起来。实测一个 6 层 Transformer 编码器encoder 部分整体包进 checkpoint 后显存占用从 11.2GB 降到 6.8GB训练步耗时从 0.43s 涨到 0.55s。如果你的显存瓶颈卡在 Activation 而不是参数checkpointing 应该是第一个尝试的方案。4.2 梯度累积绕过显存限制的 batch size 倍增方案梯度累积不是从 PyTorch 内部减少 Activation而是从训练策略层面绕开单次前向的大显存需求。假如你想要的等效 batch size 是 64但显卡一次只能塞下 batch size 16那就可以连续跑 4 个 micro-batch每步只做 forwardbackward但不立刻更新参数把梯度累积起来攒够 4 次后再optimizer.step()。实现模式很简单for i, batch in enumerate(dataloader): loss model(batch) loss loss / accumulation_steps # 归一化非常重要 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里有个很多人踩过的坑如果不对 loss 除以累积步数等效 batch size 变大后梯度幅值也会成倍增长学习率需要等比调整否则训练极易发散。另外注意 BNBatch Normalization在梯度累积下的行为差异——每个 micro-batch 单独统计均值和方差相当于用了更小的 mini-batch 做归一化某些任务上精度会有微妙变化。4.3 混合精度训练把 Activation 的存储砍一半既然 Activation 默认是 FP32 存储的那最直接的优化就是让它用 FP16 存储。PyTorch 自带的自动混合精度AMP机制在 PyTorch 1.6 之后已经相当成熟用法也简单scaler torch.cuda.amp.GradScaler() for batch in dataloader: with torch.cuda.amp.autocast(): loss model(batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()AMP 的原理是同时对三个维度做了优化算子精度前向中的 Conv、Linear 等算子自动用 FP16 计算速度明显提升。Activation 存储中间激活值以前向内算子输入输出的实际精度存储FP16 直接砍半显存。梯度缩放为了防止 FP16 下梯度过小被冲刷掉用 GradScaler 对 loss 做缩放反向传播后再缩放回原梯度。如果你用的是 NVIDIA 30 系及以上显卡AMP 支持得已经相当成熟。实测 ResNet-50 在 AMP 模式下显存减少约 40%训练速度还能提升 50% 以上。不过我见过不少同学以为 AMP 只需要加一行autocast()就万事大吉忽略了 GradScaler结果模型 loss 变成 NaN 或者不收敛的这一点值得反复强调。混合精度训练后如果你的模型里还有 BatchNorm有可能遇到方差偏大引发的数值不稳定。可以把torch.cuda.amp.autocast(enabledTrue, dtypetorch.bfloat16)改成 BF16 试试BF16 跟 FP16 的区别在于指数位更多动态范围跟 FP32 接近省去了 GradScaler 的动态调整麻烦。代价是某些卡如 A100 之后的高端卡对 BF16 的吞吐支持才最好。4.4 其他值得动手的小优化inplace 操作与激活函数替换除了以上三大手段日常代码里还有一些微观层面的显存优化适合在项目后期做细粒度调优。第一类是尽量使用inplace操作。例如nn.ReLU(inplaceTrue)可以让输出直接覆盖输入张量省去一块独立输出的显存。但注意inplace不能滥用如果该张量同时是反向传播需要的前向输入覆盖操作会导致梯度计算读取到错误的数据。PyTorch 官方文档明确标注了哪些算子的 inplace 版本安全建议只对明确无风险的位置开。第二类是激活函数替换。比如用SiLU也叫 Swish替换部分GELU在 Transformer 类模型上能省一些中间存储。某些场景下用torch.nn.functional.gelu(..., approximatetanh)也能拿到轻微收益但三层以上的复合结构不一定划算建议先 profiler 再动手。第三类是小心.detach()与切片操作的隐性拷贝。比如在自定义 Dataset 里对.numpy()的图和数据做切片某些隐藏操作会额外复制张量在显存敏感的循环中积累出不小的额外开销。真的在意显存时用torch.utils.cpp_extension或者单独算子替换一次大张量操作往往比零碎调 inplace 更为高效。5. 实测排查从显存曲线定位到具体的 Activation 位置优化做完不等于结束。我习惯在每次训练脚本跑通之后用 PyTorch 官方 profiling 工具对显存做一个现场回放确认每个模板化优化的收益也避免把宝贵的精力投在收益最低的位置。5.1 用 memory_profiler 与 torch.profiler 拿到按张量维度的显存报告步骤一先给验证代码包上torch.profilerimport torch from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue) as prof: output model(dummy_input) loss output.sum() loss.backward() print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit20))这份表格会列出每个算子自身的 CUDA 显存分配情况能够非常直观地看到你模型里哪个算子通常是某个大的矩阵乘法转置、某个残差连接的加法把显存拉爆了。数据维度越大的地方通常 Activation 也越大这时候就应该结合前文提到的方法去规划到底是 checkpointing、混合精度还是调小 batch size。步骤二我对同一个实验做了一次干预前后的纵向对比。优化手段batch size显存峰值单步耗时等效吞吐基线FP321611.2GB0.43s37.2 batch/s AMP166.8GB0.31s51.6 batch/s CheckpointingTransformer 段164.5GB0.42s38.1 batch/s 梯度累积 x2324.5GB0.84s38.1 batch/s注意表格里最后一行梯度累积并没有直接提升单次处理的吞吐它解决的是你根本没有足够显存跑等效 batch 32的问题。在你显存不足且不可换卡的前提下组合使用这三种手段才能把模型塞进有限的硬件里。5.2 常见显存误判到底是不是 Activation 吃掉了显存排查显存问题时我总结了一个快速判断经验如果你在某个位置发现显存高先判断数据是 Activation、梯度、参数、优化器状态还是 CUDA context 碎片。区分方法很简单把所有输入和参数都requires_gradFalse再前向一次看显存是否下降。如果不降问题八成不在 Activation而在数据加载、缓存池或 CUDA context 初始化上。关于 CUDA context 和显卡驱动配套问题想多说一句很多时候换了一张新显卡或者换了 PyTorch 版本跑训练时系统提示PyTorch 与驱动版本不匹配这不一定就是显存不够。正确的做法是先确认 CUDA driver、CUDA runtime 和 PyTorch 三者的版本配套关系再考虑显存优化。否则你在优化代码真正的瓶颈却是显存驱动层没走对。6. 进阶扩展torch.compile 与 FSDP 对显存生命周期的改变如果你把前面几章的方法都吃透了还想继续压榨显存效率那有两样新东西值得了解torch.compile和 FSDPFully Sharded Data Parallel。它们不是在同一个层面优化而是从编译期和分布式层面重塑了显存生命周期的分配模式。6.1 torch.compile 的算子融合效应PyTorch 2.x 已经支持torch.compile后端通过 Triton 或 Inductor 对计算图做算子融合。它能把前向传播中多个相邻小算子融合成一个融合算子执行好处有两个一是算子调度次数减少二是融合算子的 intermediate tensor 不需要在显存中落地很多情况下可以降低 Activation 的峰值占用。我自己在一个 8 层 Transformer 的实验中对比过不开 compile 时前向峰值显存 4.2GB跑torch.compile(model, modereduce-overhead)后显存峰值降到 3.5GB耗时也有 8% 左右的下降。但需要留意torch.compile与动态图的天然冲突如果模型结构中包含大量数据相关的分支编译优化可能不奏效甚至因为编译缓存管理额外占用显存。6.2 FSDP 的显存切分与 recompute 策略当单个 GPU 装不下完整模型参数和梯度时FSDP 会把参数、梯度和优化器状态给打散到多个 GPU 上。这种切分策略对显存生命周期的改变非常大每张卡不再保留完整的整层参数而是在前向用时通过 all-gather 汇聚参数算完立刻释放反向时再次汇聚。如果 FSDP 和 Activation Checkpointing 同时开启你还需要注意它们的交互策略。FSDP 的forward_prefetch和backward_prefetch参数决定了参数预抓取的时机设置不当会导致某个 GPU 的显存峰值出现在前向和反向并行双抓的时刻。经验是在显存压力大的卡上关掉backward_prefetchTrue改为False反而能降低峰值。6.3 显存调优的优先级清单最后给一份按照收益与成本排序的调优优先级建议方便你接到一个新任务时快速定位先跑通脚本用torch.cuda.memory_summary()获取基准显存水位。开启 AMP混合精度收益最高、改动最小。引入 Activation Checkpointing覆盖计算图中显存占比大的区域。再做梯度累积解决等效 batch size 不足的问题。如果参数和梯度本身成了瓶颈考虑 FSDP 或优化器状态压缩如 Adafactor 替代 Adam。最后检查数据加载管线防止 CPU 侧同步阻塞导致 GPU 等待而非显存不足。我在实际项目里按照这个优先级操作后几乎很少遇到必须换卡的情况。很多看起来很吓人的 OOM其实根本原因只是 Activation 的存储精度和生命周期没管理好。这个 PyTorch 计算图与 Autograd 的显存生命周期话题越深入越觉得有趣动态 DAG 的灵活性换来的是运行时构建开销Autograd 的自动化换来的是显存的中间缓存堆积Activation 的优化则是在时间、显存、代码复杂度之间取平衡。没有一劳永逸的银弹但了解这几层机制各自的取舍之后你就能在遇到问题时有方向感地动手尝试而不是靠瞎猜和盲试。 PyTorch动态计算图与Autograd：从DAG构建到显存优化实践
1. 动态计算图PyTorch 为什么偏要在运行时现搭 DAG不少同学从 TensorFlow 1.x 时代转过来用 PyTorch 时最不适应的就是怎么没看到图结构在 TensorFlow 1.x 里你得先用tf.Graph定义好一张静态计算图再丢给Session去跑而 PyTorch 上手就是直接写张量运算看起来跟写普通 Python 脚本几乎没区别。直到你第一次打印一个张量的grad_fn属性才发现奥妙藏在背后。1.1 计算图不是画出来的是记下来的静态计算图的构建方式是先搭图后执行图结构一旦定义基本固定改动特别麻烦。PyTorch 走的完全是另一条路它在每次前向传播执行张量运算时动态记录操作历史自动把算子之间依赖关系编织成一张有向无环图DAG。这张图不是预先画好存在某个对象里的而是通过一个个 Tensor 的grad_fn属性层层串联起来的。举个例子我用最朴素的三行代码演示这一点import torch x torch.tensor([2.0], requires_gradTrue) z x ** 2 # z x^2 y 3 * z 5 # y 3x^2 5 print(x.grad_fn) # None因为 x 是叶子节点 print(z.grad_fn) # PowBackward0 object at 0x... print(y.grad_fn) # AddBackward0 object at 0x...z.grad_fn指向的是生成 z 的那个操作的逆向函数入口。这个对象内部还保留着指向输入张量x的引用以及操作所需的元数据。PyTorch 就是靠这种链式结构把整个前向传播过程组织成一张可以从输出一路回溯到输入的图。我把动态图和静态图做了个直观对比对比维度动态图PyTorch静态图TF1.x图构建时机每次前向传播运行时动态构建先定义后执行调试体验可以用 print、断点、if/else 原生 Python 控制流需要把控制流改成图节点调试困难性能优化空间编译优化难但灵活性极高可做图优化、算子融合典型适用场景研究、动态结构模型、RL生产部署、静态推理为什么 PyTorch 选择了磁带式自动微分这条路线核心原因是研究和实验中模型结构经常需要调整网络层数、分支结构甚至损失函数都可能前一天才改过。动态图意味着你完全可以用if语句判断输入长度来决定走哪个分支因为图就是跟着 Python 执行轨迹现搭的不需要为难你。1.2 磁带式记录DAG 的建立时机和内部关联PyTorch 的自动微分原理其实可以类比成录音机或者磁带你在前向传播过程中做的每一步张量运算都像在磁带上录下操作类型、输入张量、输出张量这三要素。反向传播时自动微分系统把磁带倒过来放沿着记录倒序执行每个操作对应的求导规则把梯度逐级传回去。再来看看 DAG 在内存里到底是什么。执行一段前向传播代码时PyTorch 会在每个需要梯度的张量之间建立边关系。拿最常见的线性层举例x torch.randn(64, 512, requires_gradTrue) w torch.randn(512, 256, requires_gradTrue) b torch.randn(256, requires_gradTrue) h x w # 矩阵乘 h h b # 广播加法 out h.tanh() # 激活函数这四行代码运行时内存中形成的 DAG 大致是这样的结构x(requires_gradTrue) w(requires_gradTrue) b(requires_gradTrue) \ / / \ / / \ / / \ / / \ / / MatMulBackward / \ / \ / \ / AddBackward / \ / \ / \ / TanhBackward | out每个反向函数对象MatMulBackward、AddBackward、TanhBackward内部都保存了前向计算时的输入。这其实是一个非常重要的细节DAG 不仅保存了结构还保存了前向运算的中间结果。正因如此反向传播时才能直接复用前向的值来算梯度比如tanh激活函数的梯度需要用到前向的out值本身。这些被保存下来的中间张量正是 Activation 显存开销的根源。1.3 叶子节点与非叶子节点的命运差异凡是requires_gradTrue且由用户直接创建、不依赖其他张量的 Tensor被称为叶子节点Leaf Tensor。非叶子节点则是由运算生成的中间张量。它们在反向传播结束后的梯度处理上待遇完全不同叶子节点反向传播完成后梯度会累积在.grad属性里留给优化器更新参数使用。如果你连续调用两次loss.backward()叶子节点的梯度是叠加的所以训练循环里每步都得手动optimizer.zero_grad()。非叶子节点反向传播完梯度只作为中间传递的桥梁用完后如果不做特殊保存会直接被丢弃释放。你要是试图打印中间张量的.grad得到的往往是None。非叶子节点如果想保留梯度需要调用.retain_grad()方法提前声明。这个机制我刚开始学 PyTorch 时经常忽略直到有一天我想可视化某层特征图的梯度打印出来全是None才意识到 PyTorch 为了省内存放弃了自动保存非叶子节点梯度的能力。2. Autograd 反向求导机制链式法则如何在 DAG 上传播反向传播表面看只是调用一下.backward()但这一行调用背后隐藏着一套完整的任务调度系统从输出节点出发按照 DAG 的拓扑逆序遍历所有反向节点依次执行局部链式求导把梯度向叶子节点方向逐级传递。2.1 从链式法则到梯度传播路径复杂模型的链式法则本质上是分段求导逐级相乘。PyTorch 对每个算子都预先注册了局部求导规则比如乘法的反向规则、卷积的反向规则。反向传播时Autograd 引擎只需要把上游传来的梯度与当前算子的局部雅可比矩阵相乘就得到这个算子各输入的梯度贡献。拿最基础的复合函数y f(g(x))举例dy/dx dy/du * du/dx其中u g(x)。在 PyTorch 中这一过程被拆成两次反向节点的传播x torch.tensor([3.0], requires_gradTrue) u 2 * x 1 # u g(x)grad_fnAddBackward y u ** 2 # y f(u)grad_fnPowBackward y.backward() print(x.grad) # 12.0即 dy/du6, du/dx2, 乘积 12y.backward()执行时Autograd 引擎首先计算dy/du也就是 2u 6然后把梯度 6 传给u的反向节点AddBackward再计算du/dx 2二者相乘得到 12写入x.grad。这就是链式法则在 DAG 上传播的完整路径。2.2 Autograd 引擎的任务调度策略PyTorch 的 Autograd 引擎在反向传播时采用了基于线程池的异步执行策略。它会从输出节点开始把无依赖的梯度计算任务放入队列多个工作线程并行处理同一层级的梯度任务当一个节点的所有输入梯度都计算完成之后才能触发该节点的梯度计算。这样一层一层解锁下去直到所有叶子节点的梯度都计算完毕。这也是为什么在反向传播过程中我们有时会看到 GPU 利用率上的锯齿波动——不同层因为计算量差异在各个阶段的表现节奏天然不一样而 Autograd 引擎是动态调度的哪里的依赖先满足就先算哪里。老版本 PyTorch 在复杂模型上有时会因为线程调度的开销卡顿1.9 之后引入了一种基于链表结构的新调度实现显著降低了反向传播的线程同步开销。实测在包含大量小算子的模型上比如密集的 attention 结构新老版本的 backward 时间能差出 10%~20%。2.3 retain_graph、detach 与梯度累积的关系调用.backward()后Autograd 引擎会把整张计算图的中间缓存全部释放。如果对同一个 loss 连续调用两次.backward()比如某些 GAN 训练代码里 generator 和 discriminator 分开更新会直接报错RuntimeError: Trying to backward through the graph a second time。解决方式有两个方向在backward()里传入retain_graphTrue强制保留计算图缓存但这样内存开销和第一次前向持平。如果你的计算图很大这一招会让显存雪上加霜。更推荐的做法是搞明白梯度累积的语义如果只是想把多批数据的梯度加到一起完全可以用.detach()切断分支或者分多次调用 backward 而每次使用不同的子图输出。比如 GAN 里 discriminator 的 loss 和 generator 的 loss 通常来自不同的计算路径只需要保证共享部分参数的计算图在当前 batch 内只 backward 一次即可。还有一个高频踩坑点把 loss 的 item 标量值拿来计算后又参与梯度计算导致图被意外保留。正确做法就是把数值统计的部分.detach()出来不要让无关的 Python 运算串进计算图。3. 显存生命周期前向、反向、优化器更新显存究竟装了什么很多人用 PyTorch 训练时只看一个整体显存占用数字其实训练一步过程中显存是被不同角色瓜分的模型参数、优化器动量、前向中间激活、反向梯度、CUDA context 和计算临时缓存。搞清楚每个时刻显存里都有谁是显存优化的重要基础。3.1 一个典型训练步骤中的显存阶段变化以最常见的一个 step 为例假设你用了一张 12GB 的显卡训练 ResNet-50输入 batch size 为 32。我用torch.cuda.memory_allocated()在两个关键时点打点观测看到的现象基本如下model torchvision.models.resnet50().cuda() optimizer torch.optim.SGD(model.parameters(), lr0.01) dummy_input torch.randn(32, 3, 224, 224).cuda() torch.cuda.synchronize() print(初始显存占用:, torch.cuda.memory_allocated() / 1024**2, MB) output model(dummy_input) loss output.sum() torch.cuda.synchronize() print(前向完成显存:, torch.cuda.memory_allocated() / 1024**2, MB) loss.backward() torch.cuda.synchronize() print(反向完成显存:, torch.cuda.memory_allocated() / 1024**2, MB) optimizer.step() torch.cuda.synchronize() print(优化器更新后:, torch.cuda.memory_allocated() / 1024**2, MB)实测规律非常稳定初始显存只包含模型参数和 CUDA context大概 100MB 到 400MB 的数量级。前向完成显存猛增到几个 GB多出来的全部是各层输出的 Activation 中间激活值。反向完成相比前向完成时没有明显变化甚至可能小幅度下降因为反向计算过程中读取并释放了部分 Activation 缓存但同时产生了各参数的梯度。优化器更新后显存基本保持不变参数梯度不会因为optimizer.step()而释放除非手动清空因为梯度要保留到下一个zero_grad()。这一个周期就把显存生命周期分得明明白白参数和梯度占的只是冰山一角前向传播时堆积起来的 Activation 才是显存水位暴涨的主导因素。3.2 CUDA 缓存分配器的假释放现象还有一个让很多人迷惑的现象你用nvidia-smi看显存占用明明很高但代码里torch.cuda.memory_allocated()显示占用很低。这是因为 PyTorch 为了减少频繁调用 CUDA 底层显存分配接口的系统开销默认使用缓存分配器。Tensor 被 Python 端释放后显存块并不会真的还给显存驱动而是留在 PyTorch 的缓存池里等待复用。这个设计的代价是你在nvidia-smi中看到的显存剩余其实不准确。有一次我跑一个需要动态 batch size 的推理服务每个请求的张量形状都在变化导致缓存池里碎片化严重发生了即使memory_allocated不高但依然 OOM 的情况最后只能通过PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True缓解。顺带提一个排查技巧出现 OOM 时先运行torch.cuda.memory_summary()它会给出当前显存池里缓存块的大小分布和分配器行为明细远比nvidia-smi的信息有诊断价值。3.3 梯度传播结束后的释放时机和内存复用反向传播结束后非叶子节点的 Activation 缓存会立刻被释放回缓存池这些显存块可以马上被下一轮前向传播复用这就是为什么连续多步训练时显存占用基本持平而不会持续上涨。反过来说如果某个中间结果因为被外部变量引用而迟迟得不到释放显存占用就会异常增高。这里有三个高频错误optimizer.zero_grad()放在 backward 前面而不是每步更新后调用导致梯度叠加。在循环里把每个 batch 的输出 append 到一个 Python list计算图一直没释放list 越界增长显存逐步耗尽。开了torch.no_grad()还挂requires_gradTrue张量做验证循环虽然不会建图但历史张量被保留也会累积。4. Activation 显存优化从 checkpointing 到混合精度的完整手段既然 Activation 是显存水位的绝对主导者那如何压低它就成了训练大规模模型的必修课。业内最核心的思路说白了只有两个方向减少同时保留在显存里的中间激活或者降低每个激活的存储精度。4.1 Activation Checkpointing用计算换显存Activation Checkpointing也叫梯度检查点的原理非常直白在反向传播时本来需要前向保存的每个中间结果现在只保留若干个检查点位置的输出其余中间结果是反向传播时再临时重新前向计算的。PyTorch 提供了一个现成的 APIfrom torch.utils.checkpoint import checkpoint def transformer_block(x): # 一些计算密集型结构 return x x checkpoint(transformer_block, x)这里需要注意checkpoint 不是免费午餐。它额外触发一次前向计算所以训练时间通常会增加 20%~30%。但对于那些显存已经告急的大模型这个时间换空间的买卖非常划算可以让原本完全跑不动的模型训练起来。实测一个 6 层 Transformer 编码器encoder 部分整体包进 checkpoint 后显存占用从 11.2GB 降到 6.8GB训练步耗时从 0.43s 涨到 0.55s。如果你的显存瓶颈卡在 Activation 而不是参数checkpointing 应该是第一个尝试的方案。4.2 梯度累积绕过显存限制的 batch size 倍增方案梯度累积不是从 PyTorch 内部减少 Activation而是从训练策略层面绕开单次前向的大显存需求。假如你想要的等效 batch size 是 64但显卡一次只能塞下 batch size 16那就可以连续跑 4 个 micro-batch每步只做 forwardbackward但不立刻更新参数把梯度累积起来攒够 4 次后再optimizer.step()。实现模式很简单for i, batch in enumerate(dataloader): loss model(batch) loss loss / accumulation_steps # 归一化非常重要 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里有个很多人踩过的坑如果不对 loss 除以累积步数等效 batch size 变大后梯度幅值也会成倍增长学习率需要等比调整否则训练极易发散。另外注意 BNBatch Normalization在梯度累积下的行为差异——每个 micro-batch 单独统计均值和方差相当于用了更小的 mini-batch 做归一化某些任务上精度会有微妙变化。4.3 混合精度训练把 Activation 的存储砍一半既然 Activation 默认是 FP32 存储的那最直接的优化就是让它用 FP16 存储。PyTorch 自带的自动混合精度AMP机制在 PyTorch 1.6 之后已经相当成熟用法也简单scaler torch.cuda.amp.GradScaler() for batch in dataloader: with torch.cuda.amp.autocast(): loss model(batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()AMP 的原理是同时对三个维度做了优化算子精度前向中的 Conv、Linear 等算子自动用 FP16 计算速度明显提升。Activation 存储中间激活值以前向内算子输入输出的实际精度存储FP16 直接砍半显存。梯度缩放为了防止 FP16 下梯度过小被冲刷掉用 GradScaler 对 loss 做缩放反向传播后再缩放回原梯度。如果你用的是 NVIDIA 30 系及以上显卡AMP 支持得已经相当成熟。实测 ResNet-50 在 AMP 模式下显存减少约 40%训练速度还能提升 50% 以上。不过我见过不少同学以为 AMP 只需要加一行autocast()就万事大吉忽略了 GradScaler结果模型 loss 变成 NaN 或者不收敛的这一点值得反复强调。混合精度训练后如果你的模型里还有 BatchNorm有可能遇到方差偏大引发的数值不稳定。可以把torch.cuda.amp.autocast(enabledTrue, dtypetorch.bfloat16)改成 BF16 试试BF16 跟 FP16 的区别在于指数位更多动态范围跟 FP32 接近省去了 GradScaler 的动态调整麻烦。代价是某些卡如 A100 之后的高端卡对 BF16 的吞吐支持才最好。4.4 其他值得动手的小优化inplace 操作与激活函数替换除了以上三大手段日常代码里还有一些微观层面的显存优化适合在项目后期做细粒度调优。第一类是尽量使用inplace操作。例如nn.ReLU(inplaceTrue)可以让输出直接覆盖输入张量省去一块独立输出的显存。但注意inplace不能滥用如果该张量同时是反向传播需要的前向输入覆盖操作会导致梯度计算读取到错误的数据。PyTorch 官方文档明确标注了哪些算子的 inplace 版本安全建议只对明确无风险的位置开。第二类是激活函数替换。比如用SiLU也叫 Swish替换部分GELU在 Transformer 类模型上能省一些中间存储。某些场景下用torch.nn.functional.gelu(..., approximatetanh)也能拿到轻微收益但三层以上的复合结构不一定划算建议先 profiler 再动手。第三类是小心.detach()与切片操作的隐性拷贝。比如在自定义 Dataset 里对.numpy()的图和数据做切片某些隐藏操作会额外复制张量在显存敏感的循环中积累出不小的额外开销。真的在意显存时用torch.utils.cpp_extension或者单独算子替换一次大张量操作往往比零碎调 inplace 更为高效。5. 实测排查从显存曲线定位到具体的 Activation 位置优化做完不等于结束。我习惯在每次训练脚本跑通之后用 PyTorch 官方 profiling 工具对显存做一个现场回放确认每个模板化优化的收益也避免把宝贵的精力投在收益最低的位置。5.1 用 memory_profiler 与 torch.profiler 拿到按张量维度的显存报告步骤一先给验证代码包上torch.profilerimport torch from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue) as prof: output model(dummy_input) loss output.sum() loss.backward() print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit20))这份表格会列出每个算子自身的 CUDA 显存分配情况能够非常直观地看到你模型里哪个算子通常是某个大的矩阵乘法转置、某个残差连接的加法把显存拉爆了。数据维度越大的地方通常 Activation 也越大这时候就应该结合前文提到的方法去规划到底是 checkpointing、混合精度还是调小 batch size。步骤二我对同一个实验做了一次干预前后的纵向对比。优化手段batch size显存峰值单步耗时等效吞吐基线FP321611.2GB0.43s37.2 batch/s AMP166.8GB0.31s51.6 batch/s CheckpointingTransformer 段164.5GB0.42s38.1 batch/s 梯度累积 x2324.5GB0.84s38.1 batch/s注意表格里最后一行梯度累积并没有直接提升单次处理的吞吐它解决的是你根本没有足够显存跑等效 batch 32的问题。在你显存不足且不可换卡的前提下组合使用这三种手段才能把模型塞进有限的硬件里。5.2 常见显存误判到底是不是 Activation 吃掉了显存排查显存问题时我总结了一个快速判断经验如果你在某个位置发现显存高先判断数据是 Activation、梯度、参数、优化器状态还是 CUDA context 碎片。区分方法很简单把所有输入和参数都requires_gradFalse再前向一次看显存是否下降。如果不降问题八成不在 Activation而在数据加载、缓存池或 CUDA context 初始化上。关于 CUDA context 和显卡驱动配套问题想多说一句很多时候换了一张新显卡或者换了 PyTorch 版本跑训练时系统提示PyTorch 与驱动版本不匹配这不一定就是显存不够。正确的做法是先确认 CUDA driver、CUDA runtime 和 PyTorch 三者的版本配套关系再考虑显存优化。否则你在优化代码真正的瓶颈却是显存驱动层没走对。6. 进阶扩展torch.compile 与 FSDP 对显存生命周期的改变如果你把前面几章的方法都吃透了还想继续压榨显存效率那有两样新东西值得了解torch.compile和 FSDPFully Sharded Data Parallel。它们不是在同一个层面优化而是从编译期和分布式层面重塑了显存生命周期的分配模式。6.1 torch.compile 的算子融合效应PyTorch 2.x 已经支持torch.compile后端通过 Triton 或 Inductor 对计算图做算子融合。它能把前向传播中多个相邻小算子融合成一个融合算子执行好处有两个一是算子调度次数减少二是融合算子的 intermediate tensor 不需要在显存中落地很多情况下可以降低 Activation 的峰值占用。我自己在一个 8 层 Transformer 的实验中对比过不开 compile 时前向峰值显存 4.2GB跑torch.compile(model, modereduce-overhead)后显存峰值降到 3.5GB耗时也有 8% 左右的下降。但需要留意torch.compile与动态图的天然冲突如果模型结构中包含大量数据相关的分支编译优化可能不奏效甚至因为编译缓存管理额外占用显存。6.2 FSDP 的显存切分与 recompute 策略当单个 GPU 装不下完整模型参数和梯度时FSDP 会把参数、梯度和优化器状态给打散到多个 GPU 上。这种切分策略对显存生命周期的改变非常大每张卡不再保留完整的整层参数而是在前向用时通过 all-gather 汇聚参数算完立刻释放反向时再次汇聚。如果 FSDP 和 Activation Checkpointing 同时开启你还需要注意它们的交互策略。FSDP 的forward_prefetch和backward_prefetch参数决定了参数预抓取的时机设置不当会导致某个 GPU 的显存峰值出现在前向和反向并行双抓的时刻。经验是在显存压力大的卡上关掉backward_prefetchTrue改为False反而能降低峰值。6.3 显存调优的优先级清单最后给一份按照收益与成本排序的调优优先级建议方便你接到一个新任务时快速定位先跑通脚本用torch.cuda.memory_summary()获取基准显存水位。开启 AMP混合精度收益最高、改动最小。引入 Activation Checkpointing覆盖计算图中显存占比大的区域。再做梯度累积解决等效 batch size 不足的问题。如果参数和梯度本身成了瓶颈考虑 FSDP 或优化器状态压缩如 Adafactor 替代 Adam。最后检查数据加载管线防止 CPU 侧同步阻塞导致 GPU 等待而非显存不足。我在实际项目里按照这个优先级操作后几乎很少遇到必须换卡的情况。很多看起来很吓人的 OOM其实根本原因只是 Activation 的存储精度和生命周期没管理好。这个 PyTorch 计算图与 Autograd 的显存生命周期话题越深入越觉得有趣动态 DAG 的灵活性换来的是运行时构建开销Autograd 的自动化换来的是显存的中间缓存堆积Activation 的优化则是在时间、显存、代码复杂度之间取平衡。没有一劳永逸的银弹但了解这几层机制各自的取舍之后你就能在遇到问题时有方向感地动手尝试而不是靠瞎猜和盲试。