很多人以为跨过 PyTorch 安装、跑通一个 MNIST 训练脚本就等于入门了结果一换真实数据集、一上大模型立刻被两个问题砸懵一个是loss.backward()之后某个参数梯度是 None另一个是 batch 稍微调大一点就CUDA out of memory。这两个问题看似毫无关联根子其实都在同一处——你对 PyTorch 底层那张计算图的理解以及 Autograd 分配显存、释放显存的节奏完全没概念。这一讲我就把这块硬骨头拆开揉碎从动态 DAG 的构建、反向求导的执行流程一直讲到 Activation 在整个训练生命周期里怎么吃掉你的显存以及对应的优化手段。我尽量不堆公式用实际代码和训练中的真实场景来讲。如果你已经能正常运行训练脚本但总感觉 PyTorch 是个黑盒这篇就是给你捅破窗户纸用的。1. 动态计算图PyTorch 的训练基石是怎么搭建的1.1 边执行边建图的 DAG 到底是什么意思PyTorch 和早期 TensorFlow 最本质的区别在于一个是动态图、一个是静态图。静态图是你先画一张完整的图纸再照着图纸施工动态图是边走边画每走一步就把这一步的依赖关系记下来。这个依赖关系就是一张有向无环图DAG。你在 Python 里写y x 1这行代码在执行的那一刻PyTorch 的 Autograd 引擎就悄悄记录了一条边x - y。x和y都是节点1这个操作是连接它们的边。因为整个图的构建跟着 Python 的控制流走所以叫动态。这里最容易忽略的关键点是并不是所有张量运算都会被记录到图里。只有requires_gradTrue的张量或者由这类张量通过运算得到的新张量才会被纳入跟踪范围。你训练时喂进去的输入数据通常不需要梯度但模型参数默认就是requires_gradTrue所以只要模型参数参与了运算这条链路就自动进了图。我自己的理解是把它想象成一次做菜过程用到的每一种食材是节点切洋葱、下锅炒这些动作是边。图记录的不是最终的菜而是这道菜做过哪些步骤、每一步用了什么。反向传播的时候你只需要沿着这条记录反向走一遍就能算出每一步的火候该往哪个方向调。1.2 Tensor 三件套data、grad 和 grad_fn一个被 Autograd 跟踪的张量身上至少有三个值得关心的属性data存数值grad存梯度grad_fn存产生这个张量的运算函数。grad_fn就是你在 debug 时最该看的字段它帮你追溯这个张量是怎么来的。import torch x torch.randn(4, 8, requires_gradTrue) w torch.randn(8, 1, requires_gradTrue) b torch.randn(1, requires_gradTrue) y x w b loss y.pow(2).mean() print(loss.grad_fn) # MeanBackward0 object at ... print(y.grad_fn) # AddBackward0 object at ... print(x.grad_fn) # None为什么x.grad_fn是 None因为它是一个叶子张量Leaf Tensor是用户直接创建的不是某个运算的产物。在反向传播时只有叶子张量会被累积梯度非叶子节点的梯度默认用完就释放这是 Autograd 显存策略里非常关键的一点后面讲显存生命周期还会提到。顺着grad_fn你可以一直追溯到最源头。每个grad_fn内部还有next_functions指向它依赖的输入节点。比如y.grad_fnAddBackward的next_functions里就包含了x w的结果和b。这意味着你可以从loss开始把整张计算图完整地走一遍不需要任何额外的存储结构。2. Autograd 反向求导机制梯度是怎么顺着 DAG 流回去的2.1 链式法则在计算图上的实际执行路径反向传播的数学本质是微积分里的链式法则。假设有这样一个链路x - h - y - loss那∂loss/∂x ∂loss/∂y · ∂y/∂h · ∂h/∂x。Autograd 做的事情就是沿着 DAG 的逆拓扑序把这个乘积一路累积下来。很多教程到这一步就结束了但实际训练中你会碰到更复杂的情况一个节点被下游多个节点使用。比如z a b之后z同时参与了p z * 2和q z - 1。那么反向传播时z的梯度应该是p传来的梯度和q传来的梯度相加。这就是为什么梯度是累积的——多个分支的梯度汇合到一个节点上必须求和。实现层面PyTorch 的loss.backward()会触发一次从loss开始的逆序遍历。每个grad_fn知道自己输入输出之间的局部梯度local gradient它接收到上游传来的梯度后乘以局部梯度再继续往下游传。这个局部梯度不需要你手动算PyTorch 对每个算子都内置了对应的反向函数。你常见的AddBackward、MulBackward、ReluBackward这些类就是干这个的。我在实践中发现一个特别有用的调试技巧用torch.autograd.gradcheck验证自定义算子的梯度是否正确用torch.autograd.set_detect_anomaly(True)来定位 NaN 梯度到底出现在哪一层。前者适合写自定义扩展时用后者适合训练突然出现 NaN 时用。不过set_detect_anomaly会拖慢训练速度我只会在调试时临时打开排查完立刻关掉。2.2 backward() 的完整执行流程与三个容易踩坑的 API一次标准的loss.backward()背后PyTorch 大致做了这几件事先检查计算图是否存在且可导然后从loss出发沿反向拓扑序遍历所有参与的grad_fn每经过一个节点计算局部梯度并乘以上游传来的梯度最终把梯度累积到叶子张量的.grad属性里。这里有几个每个 PyTorch 用户都该刻在脑子的细节第一梯度默认是累加的不会自动清零。你连续调用三次loss.backward()同一个参数的.grad会累计三份梯度。所以如果训练循环里忘了optimizer.zero_grad()loss 会像脱缰野马一样乱跳。这也解释了为什么梯度累积Gradient Accumulation技术能实现等价大 batch——每步不清零攒够 N 步再更新一次参数数学上等同于 batch 扩大 N 倍。第二.detach()和torch.no_grad()的边界完全不同。.detach()是从当前计算图中剪断连接返回一个不跟踪梯度的新张量但后续如果让这个新张量参与新运算新运算可以重新开一条新的记录链。torch.no_grad()则是给整个上下文环境加了一个免跟踪的开关常用于推理阶段和计算验证指标时。如果你在训练过程里想拿一个中间结果算个指标、又不想让它影响梯度用.detach()更精准。第三requires_grad_()可以动态切换叶子张量的跟踪状态。我试过在迁移学习中只想微调分类头、冻结特征提取层。如果不加处理地让整个网络参与反向显存和训练时间都会白白浪费。正确做法是把该冻结的层参数param.requires_grad_(False)同时优化器只传入需要更新的参数。这样 Autograd 构建计算图时被冻结的参数就不会进入图反向自然就绕过了它们。第四inplace 操作是 Autograd 的天敌。你在requires_gradTrue的叶子张量上做x.add_(1)几乎必然会报错a leaf Variable that requires grad is being used in an in-place operation。原因是 inplace 操作会覆盖前向时保存的数值而反向传播需要用到原始值计算局部梯度。就算不报错也很容易得到错误的梯度。所以我训练模型时激活函数能不用inplaceTrue就不用尤其是网络第一层附近避免踩雷。3. Activation 显存生命周期训练时显存到底被谁吃掉了3.1 一张图看清训练时的显存分配很多初学者以为显存就是拿来放模型参数的这完全是误解。训练状态下显存里同时躺着四类东西模型参数、参数梯度、优化器状态、前向过程中的 Activation激活值。模型参数数量固定占用的显存固定。参数梯度和参数同等大小反向传播时分配。优化器状态这玩意儿最容易被忽略。SGD 没有额外状态但 Adam 要保存一阶动量m和二阶动量v每个参数对应两份额外副本等于优化器状态是参数量的两倍。Activation包括每一层的输入、输出、中间缓冲是训练过程显存消耗的大头。拿 ResNet-50 举例输入 224×224、batch size 32模型参数约 2500 万FP32 下参数本身大约 100MB梯度和 Adam 状态再加 300MB。但前向传播逐层计算时每一层都要保存输入特征图几十上百层的特征图累加起来动辄 500MB 到 1GB 以上。分辨率越大、batch 越大、通道数越多Activation 的膨胀越夸张。3.2 前向保存、反向释放的动态过程Activation 显存最核心的特征是生命周期非常短但又必须精准存活到反向用到它的那一刻。前向传播时每经过一层这层的输入 Activation 就被记录到计算图中等待反向时使用所以前向过程显存呈单调增长。到了前向结束所有 Activation 都存着这是整个训练迭代中显存占用最高的时刻。反向传播时Autograd 从最后一层往前计算梯度。每算完一层的梯度该层对应的 Activation 就完成了历史使命可以被释放。所以反向过程中显存逐渐下降但也伴随产生参数梯度、中间梯度等新开销。整个训练迭代的显存峰值基本就出现在前向刚结束、反向刚开始的那个瞬间。我建议你在自己的代码里跑一下这个显存监控片段比看任何虚拟机理论都直观import torch model torchvision.models.resnet50().cuda() inputs torch.randn(32, 3, 224, 224).cuda() labels torch.randint(0, 1000, (32,)).cuda() optimizer torch.optim.Adam(model.parameters()) criterion torch.nn.CrossEntropyLoss() torch.cuda.reset_peak_memory_stats() outputs model(inputs) # 前向结束显存接近峰值 peak_after_forward torch.cuda.max_memory_allocated() loss criterion(outputs, labels) loss.backward() # 反向结束参数梯度已就位 peak_after_backward torch.cuda.max_memory_allocated() print(f前向后峰值: {peak_after_forward / 1024**2:.2f} MB) print(f反向后峰值: {peak_after_backward / 1024**2:.2f} MB)你会发现peak_after_forward和peak_after_backward通常非常接近甚至一模一样因为max_memory_allocated记录的是历史峰值。把reset_peak_memory_stats()放在反向之后再调用一次才能看到参数梯度的增量开销。4. Activation 显存优化实战从 checkpointing 到梯度累积4.1 Activation Checkpointing用计算量换显存理解了 Activation 的生命周期你就明白一个道理如果前向传播时不保存某些层中间结果而是在反向需要时重新计算一遍显存峰值就能显著下降。这个策略就叫 Activation Checkpointing也叫梯度检查点Gradient Checkpointing。具体做法是把网络切成若干段段与段之间的边界保存输出段内部不保存中间 Activation。反向传播时某一需要梯度经过这个段就临时把段内前向重算一遍用重算出来的结果计算梯度。代价是训练时间大约增加 30%~50%但显存峰值能砍掉一半甚至更多。PyTorch 内置了非常便捷的接口from torch.utils.checkpoint import checkpoint class CheckpointedBlock(nn.Module): def __init__(self, block): super().__init__() self.block block def forward(self, x): # use_reentrantFalse 是 PyTorch 2.x 推荐写法 return checkpoint(self.block, x, use_reentrantFalse)我把大型 ResNet、Transformer 的每个 stage 或每 N 层包在checkpoint里其他层保持正常工作流。实际用的判断标准是显存吃紧但训练时间还能接受就多包几层显存充足就不开因为 checkpoint 毕竟有重算开销没必要白白浪费时间。有一个坑必须提醒被checkpoint包裹的模块里如果有 BatchNorm 层要特别注意当前 PyTorch 版本对 BatchNorm 统计量的处理。某些版本在重算模式下统计量会不一致导致验证指标和训练指标出现诡异差异。我的应对方案是包含 BatchNorm 的模型我会优先考虑减小 batch size 或开 AMP而不是用 checkpointing如果用就选不含 BN 的模型结构比如 Transformer。4.2 梯度累积与混合精度降低显存峰值的另外两条路梯度累积是另一种降低单次迭代显存的方式同一批数据拆成多个 micro-batch分别前向和反向但不清零梯度攒够若干个 micro-batch 的梯度后再更新一次参数。等价于把一个大 batch 拆成好几份每次只用小 batch 的显存量。accum_steps 4 scaler torch.cuda.amp.GradScaler() for step, (inputs, labels) in enumerate(train_loader): with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) / accum_steps # 关键是平均 scaler.scale(loss).backward() if (step 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()注意这里我做了两件事将 loss 除以accum_steps让累积 N 次的梯度等价于原始 batch 平均梯度把原先的 loss 本身也放在autocast里让模型前向全流程使用 FP16 计算。这个scaler是混合精度训练AMP的标准组件因为 FP16 下梯度容易下溢需要动态缩放到一个安全范围step时再缩放回去并更新参数。AMP 对显存的帮助是实打实的模型参数和 Activation 以 FP16 保存内存直接减半。和 gradient checkpointing 组合使用很多原来只能跑 batch 4 的模型能跑到 batch 16 甚至 32。我自己在 8GB 显存的旧卡上训练一个中等规模 Transformer就是靠AMP 梯度累积 checkpointing三件套硬生生把 batch size 提了上去训练效果比被迫用小 batch 明显稳定几个量级。5. 显存问题排查与 Autograd 常见坑5.1 CUDA out of memory 的标准排查顺序OOM 是最常见的报错但很多人一看到它就直接调小 batch。我的排查顺序永远是先看显存到底用在哪里再决定动作。先把torch.cuda.max_memory_allocated()打印出来看看真实分配情况。然后用torch.cuda.memory_summary()查看更细的内存分配统计。我遇到过的情况里至少有三成不是 Activation 爆了而是前一版代码残留的内存没释放干净。在 PyTorch 中退出训练循环后显存不会立刻释放需要配合torch.cuda.empty_cache()手动释放缓存块或者干脆重启进程。排查一圈之后按成本从低到高的顺序排列优化手段现象优先手段说明整个训练阶段 OOM减小 batch size最直接也最无脑的办法副作用小降低 batch 后仍然 OOM开启 AMP半精度计算显存直接减半AMP 后仍然 OOM梯度累积用小 batch 等价大 batch训练有效性不降前向峰值仍然太高Activation checkpointing用时间换空间重算代价可控验证阶段 OOM检查是否用了torch.no_grad()很多推理显存浪费来自忘记关梯度跟踪还有一个鲜为人知的技巧如果你用的是DataLoadernum_workers开太大也会导致系统整体内存吃紧进而拖累显存分配效率。我跑大尺寸图片时会把num_workers从默认的 0 调到 2 或 4配合persistent_workersTrue内存占用反而更平滑。5.2 Autograd 常见错误速查与独家避坑技巧这里贴一份我在实际项目里反复用到的错误排查表每一条都来自真实踩坑让我想想有没有其他坑比如retain_graph第二次 backward 报错时要加 retain_graphTrue但会显著增加显存适合只在特殊场景使用。报错或现象根因处理方式某参数.grad是 None参数没参与当前 batch 计算或 requires_gradFalse检查模型前向路径确认参数真的用到了a leaf Variable ... in-place operation对叶子张量做了原地修改改用x.data或创建新张量训练中避免 inplace第二次调用 backward 报错默认计算图已被释放无法复用需要二次梯度时传retain_graphTrue用毕及时释放one of the variables needed for gradient computation has been modified by an inplace operation非叶子张量被原地修改重点查激活函数和自定义层里的 inplace 操作Expected scalar type Float but found Half混合精度下算子类型不匹配检查自定义层是否将输入转换成了 FP32或反向时没保持一致grad can be implicitly created only for scalar outputs对非标量张量直接调用 backward传入与张量形状相同的梯度权重向量或用.sum()/.mean()先聚合排查梯度问题有一个我每次都会先做的动作开启torch.autograd.set_detect_anomaly(True)运行到报错位置PyTorch 会明确告诉你到底是哪个算子、哪一行代码触发了异常。这个开关不需要永久开启只是定位问题的临时工具。最后一个独家建议不要把detach()当成逃避深思熟虑的万能钥匙。我见过有人为了修复梯度异常到处乱加.detach()最后每个 batch 的梯度都变成零模型彻底停止学习。detach()的本质是主动切断梯度流它的使用场景应该是你确定这里确实不需要回传梯度比如计算辅助的监控指标、或者把强化学习中的某项 reward 当常数处理。不确定时先找出梯度的来源再用torch.autograd.grad或打印grad_fn链去确认远比乱剪图安全。我自己的习惯是每次搭建一个新模型结构先在 CPU 上用一个小随机输入从头到尾走一遍前向和反向用torch.autograd.gradcheck校验梯度确认无误再上 GPU。这个习惯帮我省下的 debug 时间比任何优化技巧都值钱。搞懂了计算图和 Autograd 的内存节奏你会发现 PyTorch 里的显存不是玄学而是一本你随时能翻开的账本。
很多人以为跨过 PyTorch 安装、跑通一个 MNIST 训练脚本就等于入门了结果一换真实数据集、一上大模型立刻被两个问题砸懵一个是loss.backward()之后某个参数梯度是 None另一个是 batch 稍微调大一点就CUDA out of memory。这两个问题看似毫无关联根子其实都在同一处——你对 PyTorch 底层那张计算图的理解以及 Autograd 分配显存、释放显存的节奏完全没概念。这一讲我就把这块硬骨头拆开揉碎从动态 DAG 的构建、反向求导的执行流程一直讲到 Activation 在整个训练生命周期里怎么吃掉你的显存以及对应的优化手段。我尽量不堆公式用实际代码和训练中的真实场景来讲。如果你已经能正常运行训练脚本但总感觉 PyTorch 是个黑盒这篇就是给你捅破窗户纸用的。1. 动态计算图PyTorch 的训练基石是怎么搭建的1.1 边执行边建图的 DAG 到底是什么意思PyTorch 和早期 TensorFlow 最本质的区别在于一个是动态图、一个是静态图。静态图是你先画一张完整的图纸再照着图纸施工动态图是边走边画每走一步就把这一步的依赖关系记下来。这个依赖关系就是一张有向无环图DAG。你在 Python 里写y x 1这行代码在执行的那一刻PyTorch 的 Autograd 引擎就悄悄记录了一条边x - y。x和y都是节点1这个操作是连接它们的边。因为整个图的构建跟着 Python 的控制流走所以叫动态。这里最容易忽略的关键点是并不是所有张量运算都会被记录到图里。只有requires_gradTrue的张量或者由这类张量通过运算得到的新张量才会被纳入跟踪范围。你训练时喂进去的输入数据通常不需要梯度但模型参数默认就是requires_gradTrue所以只要模型参数参与了运算这条链路就自动进了图。我自己的理解是把它想象成一次做菜过程用到的每一种食材是节点切洋葱、下锅炒这些动作是边。图记录的不是最终的菜而是这道菜做过哪些步骤、每一步用了什么。反向传播的时候你只需要沿着这条记录反向走一遍就能算出每一步的火候该往哪个方向调。1.2 Tensor 三件套data、grad 和 grad_fn一个被 Autograd 跟踪的张量身上至少有三个值得关心的属性data存数值grad存梯度grad_fn存产生这个张量的运算函数。grad_fn就是你在 debug 时最该看的字段它帮你追溯这个张量是怎么来的。import torch x torch.randn(4, 8, requires_gradTrue) w torch.randn(8, 1, requires_gradTrue) b torch.randn(1, requires_gradTrue) y x w b loss y.pow(2).mean() print(loss.grad_fn) # MeanBackward0 object at ... print(y.grad_fn) # AddBackward0 object at ... print(x.grad_fn) # None为什么x.grad_fn是 None因为它是一个叶子张量Leaf Tensor是用户直接创建的不是某个运算的产物。在反向传播时只有叶子张量会被累积梯度非叶子节点的梯度默认用完就释放这是 Autograd 显存策略里非常关键的一点后面讲显存生命周期还会提到。顺着grad_fn你可以一直追溯到最源头。每个grad_fn内部还有next_functions指向它依赖的输入节点。比如y.grad_fnAddBackward的next_functions里就包含了x w的结果和b。这意味着你可以从loss开始把整张计算图完整地走一遍不需要任何额外的存储结构。2. Autograd 反向求导机制梯度是怎么顺着 DAG 流回去的2.1 链式法则在计算图上的实际执行路径反向传播的数学本质是微积分里的链式法则。假设有这样一个链路x - h - y - loss那∂loss/∂x ∂loss/∂y · ∂y/∂h · ∂h/∂x。Autograd 做的事情就是沿着 DAG 的逆拓扑序把这个乘积一路累积下来。很多教程到这一步就结束了但实际训练中你会碰到更复杂的情况一个节点被下游多个节点使用。比如z a b之后z同时参与了p z * 2和q z - 1。那么反向传播时z的梯度应该是p传来的梯度和q传来的梯度相加。这就是为什么梯度是累积的——多个分支的梯度汇合到一个节点上必须求和。实现层面PyTorch 的loss.backward()会触发一次从loss开始的逆序遍历。每个grad_fn知道自己输入输出之间的局部梯度local gradient它接收到上游传来的梯度后乘以局部梯度再继续往下游传。这个局部梯度不需要你手动算PyTorch 对每个算子都内置了对应的反向函数。你常见的AddBackward、MulBackward、ReluBackward这些类就是干这个的。我在实践中发现一个特别有用的调试技巧用torch.autograd.gradcheck验证自定义算子的梯度是否正确用torch.autograd.set_detect_anomaly(True)来定位 NaN 梯度到底出现在哪一层。前者适合写自定义扩展时用后者适合训练突然出现 NaN 时用。不过set_detect_anomaly会拖慢训练速度我只会在调试时临时打开排查完立刻关掉。2.2 backward() 的完整执行流程与三个容易踩坑的 API一次标准的loss.backward()背后PyTorch 大致做了这几件事先检查计算图是否存在且可导然后从loss出发沿反向拓扑序遍历所有参与的grad_fn每经过一个节点计算局部梯度并乘以上游传来的梯度最终把梯度累积到叶子张量的.grad属性里。这里有几个每个 PyTorch 用户都该刻在脑子的细节第一梯度默认是累加的不会自动清零。你连续调用三次loss.backward()同一个参数的.grad会累计三份梯度。所以如果训练循环里忘了optimizer.zero_grad()loss 会像脱缰野马一样乱跳。这也解释了为什么梯度累积Gradient Accumulation技术能实现等价大 batch——每步不清零攒够 N 步再更新一次参数数学上等同于 batch 扩大 N 倍。第二.detach()和torch.no_grad()的边界完全不同。.detach()是从当前计算图中剪断连接返回一个不跟踪梯度的新张量但后续如果让这个新张量参与新运算新运算可以重新开一条新的记录链。torch.no_grad()则是给整个上下文环境加了一个免跟踪的开关常用于推理阶段和计算验证指标时。如果你在训练过程里想拿一个中间结果算个指标、又不想让它影响梯度用.detach()更精准。第三requires_grad_()可以动态切换叶子张量的跟踪状态。我试过在迁移学习中只想微调分类头、冻结特征提取层。如果不加处理地让整个网络参与反向显存和训练时间都会白白浪费。正确做法是把该冻结的层参数param.requires_grad_(False)同时优化器只传入需要更新的参数。这样 Autograd 构建计算图时被冻结的参数就不会进入图反向自然就绕过了它们。第四inplace 操作是 Autograd 的天敌。你在requires_gradTrue的叶子张量上做x.add_(1)几乎必然会报错a leaf Variable that requires grad is being used in an in-place operation。原因是 inplace 操作会覆盖前向时保存的数值而反向传播需要用到原始值计算局部梯度。就算不报错也很容易得到错误的梯度。所以我训练模型时激活函数能不用inplaceTrue就不用尤其是网络第一层附近避免踩雷。3. Activation 显存生命周期训练时显存到底被谁吃掉了3.1 一张图看清训练时的显存分配很多初学者以为显存就是拿来放模型参数的这完全是误解。训练状态下显存里同时躺着四类东西模型参数、参数梯度、优化器状态、前向过程中的 Activation激活值。模型参数数量固定占用的显存固定。参数梯度和参数同等大小反向传播时分配。优化器状态这玩意儿最容易被忽略。SGD 没有额外状态但 Adam 要保存一阶动量m和二阶动量v每个参数对应两份额外副本等于优化器状态是参数量的两倍。Activation包括每一层的输入、输出、中间缓冲是训练过程显存消耗的大头。拿 ResNet-50 举例输入 224×224、batch size 32模型参数约 2500 万FP32 下参数本身大约 100MB梯度和 Adam 状态再加 300MB。但前向传播逐层计算时每一层都要保存输入特征图几十上百层的特征图累加起来动辄 500MB 到 1GB 以上。分辨率越大、batch 越大、通道数越多Activation 的膨胀越夸张。3.2 前向保存、反向释放的动态过程Activation 显存最核心的特征是生命周期非常短但又必须精准存活到反向用到它的那一刻。前向传播时每经过一层这层的输入 Activation 就被记录到计算图中等待反向时使用所以前向过程显存呈单调增长。到了前向结束所有 Activation 都存着这是整个训练迭代中显存占用最高的时刻。反向传播时Autograd 从最后一层往前计算梯度。每算完一层的梯度该层对应的 Activation 就完成了历史使命可以被释放。所以反向过程中显存逐渐下降但也伴随产生参数梯度、中间梯度等新开销。整个训练迭代的显存峰值基本就出现在前向刚结束、反向刚开始的那个瞬间。我建议你在自己的代码里跑一下这个显存监控片段比看任何虚拟机理论都直观import torch model torchvision.models.resnet50().cuda() inputs torch.randn(32, 3, 224, 224).cuda() labels torch.randint(0, 1000, (32,)).cuda() optimizer torch.optim.Adam(model.parameters()) criterion torch.nn.CrossEntropyLoss() torch.cuda.reset_peak_memory_stats() outputs model(inputs) # 前向结束显存接近峰值 peak_after_forward torch.cuda.max_memory_allocated() loss criterion(outputs, labels) loss.backward() # 反向结束参数梯度已就位 peak_after_backward torch.cuda.max_memory_allocated() print(f前向后峰值: {peak_after_forward / 1024**2:.2f} MB) print(f反向后峰值: {peak_after_backward / 1024**2:.2f} MB)你会发现peak_after_forward和peak_after_backward通常非常接近甚至一模一样因为max_memory_allocated记录的是历史峰值。把reset_peak_memory_stats()放在反向之后再调用一次才能看到参数梯度的增量开销。4. Activation 显存优化实战从 checkpointing 到梯度累积4.1 Activation Checkpointing用计算量换显存理解了 Activation 的生命周期你就明白一个道理如果前向传播时不保存某些层中间结果而是在反向需要时重新计算一遍显存峰值就能显著下降。这个策略就叫 Activation Checkpointing也叫梯度检查点Gradient Checkpointing。具体做法是把网络切成若干段段与段之间的边界保存输出段内部不保存中间 Activation。反向传播时某一需要梯度经过这个段就临时把段内前向重算一遍用重算出来的结果计算梯度。代价是训练时间大约增加 30%~50%但显存峰值能砍掉一半甚至更多。PyTorch 内置了非常便捷的接口from torch.utils.checkpoint import checkpoint class CheckpointedBlock(nn.Module): def __init__(self, block): super().__init__() self.block block def forward(self, x): # use_reentrantFalse 是 PyTorch 2.x 推荐写法 return checkpoint(self.block, x, use_reentrantFalse)我把大型 ResNet、Transformer 的每个 stage 或每 N 层包在checkpoint里其他层保持正常工作流。实际用的判断标准是显存吃紧但训练时间还能接受就多包几层显存充足就不开因为 checkpoint 毕竟有重算开销没必要白白浪费时间。有一个坑必须提醒被checkpoint包裹的模块里如果有 BatchNorm 层要特别注意当前 PyTorch 版本对 BatchNorm 统计量的处理。某些版本在重算模式下统计量会不一致导致验证指标和训练指标出现诡异差异。我的应对方案是包含 BatchNorm 的模型我会优先考虑减小 batch size 或开 AMP而不是用 checkpointing如果用就选不含 BN 的模型结构比如 Transformer。4.2 梯度累积与混合精度降低显存峰值的另外两条路梯度累积是另一种降低单次迭代显存的方式同一批数据拆成多个 micro-batch分别前向和反向但不清零梯度攒够若干个 micro-batch 的梯度后再更新一次参数。等价于把一个大 batch 拆成好几份每次只用小 batch 的显存量。accum_steps 4 scaler torch.cuda.amp.GradScaler() for step, (inputs, labels) in enumerate(train_loader): with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) / accum_steps # 关键是平均 scaler.scale(loss).backward() if (step 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()注意这里我做了两件事将 loss 除以accum_steps让累积 N 次的梯度等价于原始 batch 平均梯度把原先的 loss 本身也放在autocast里让模型前向全流程使用 FP16 计算。这个scaler是混合精度训练AMP的标准组件因为 FP16 下梯度容易下溢需要动态缩放到一个安全范围step时再缩放回去并更新参数。AMP 对显存的帮助是实打实的模型参数和 Activation 以 FP16 保存内存直接减半。和 gradient checkpointing 组合使用很多原来只能跑 batch 4 的模型能跑到 batch 16 甚至 32。我自己在 8GB 显存的旧卡上训练一个中等规模 Transformer就是靠AMP 梯度累积 checkpointing三件套硬生生把 batch size 提了上去训练效果比被迫用小 batch 明显稳定几个量级。5. 显存问题排查与 Autograd 常见坑5.1 CUDA out of memory 的标准排查顺序OOM 是最常见的报错但很多人一看到它就直接调小 batch。我的排查顺序永远是先看显存到底用在哪里再决定动作。先把torch.cuda.max_memory_allocated()打印出来看看真实分配情况。然后用torch.cuda.memory_summary()查看更细的内存分配统计。我遇到过的情况里至少有三成不是 Activation 爆了而是前一版代码残留的内存没释放干净。在 PyTorch 中退出训练循环后显存不会立刻释放需要配合torch.cuda.empty_cache()手动释放缓存块或者干脆重启进程。排查一圈之后按成本从低到高的顺序排列优化手段现象优先手段说明整个训练阶段 OOM减小 batch size最直接也最无脑的办法副作用小降低 batch 后仍然 OOM开启 AMP半精度计算显存直接减半AMP 后仍然 OOM梯度累积用小 batch 等价大 batch训练有效性不降前向峰值仍然太高Activation checkpointing用时间换空间重算代价可控验证阶段 OOM检查是否用了torch.no_grad()很多推理显存浪费来自忘记关梯度跟踪还有一个鲜为人知的技巧如果你用的是DataLoadernum_workers开太大也会导致系统整体内存吃紧进而拖累显存分配效率。我跑大尺寸图片时会把num_workers从默认的 0 调到 2 或 4配合persistent_workersTrue内存占用反而更平滑。5.2 Autograd 常见错误速查与独家避坑技巧这里贴一份我在实际项目里反复用到的错误排查表每一条都来自真实踩坑让我想想有没有其他坑比如retain_graph第二次 backward 报错时要加 retain_graphTrue但会显著增加显存适合只在特殊场景使用。报错或现象根因处理方式某参数.grad是 None参数没参与当前 batch 计算或 requires_gradFalse检查模型前向路径确认参数真的用到了a leaf Variable ... in-place operation对叶子张量做了原地修改改用x.data或创建新张量训练中避免 inplace第二次调用 backward 报错默认计算图已被释放无法复用需要二次梯度时传retain_graphTrue用毕及时释放one of the variables needed for gradient computation has been modified by an inplace operation非叶子张量被原地修改重点查激活函数和自定义层里的 inplace 操作Expected scalar type Float but found Half混合精度下算子类型不匹配检查自定义层是否将输入转换成了 FP32或反向时没保持一致grad can be implicitly created only for scalar outputs对非标量张量直接调用 backward传入与张量形状相同的梯度权重向量或用.sum()/.mean()先聚合排查梯度问题有一个我每次都会先做的动作开启torch.autograd.set_detect_anomaly(True)运行到报错位置PyTorch 会明确告诉你到底是哪个算子、哪一行代码触发了异常。这个开关不需要永久开启只是定位问题的临时工具。最后一个独家建议不要把detach()当成逃避深思熟虑的万能钥匙。我见过有人为了修复梯度异常到处乱加.detach()最后每个 batch 的梯度都变成零模型彻底停止学习。detach()的本质是主动切断梯度流它的使用场景应该是你确定这里确实不需要回传梯度比如计算辅助的监控指标、或者把强化学习中的某项 reward 当常数处理。不确定时先找出梯度的来源再用torch.autograd.grad或打印grad_fn链去确认远比乱剪图安全。我自己的习惯是每次搭建一个新模型结构先在 CPU 上用一个小随机输入从头到尾走一遍前向和反向用torch.autograd.gradcheck校验梯度确认无误再上 GPU。这个习惯帮我省下的 debug 时间比任何优化技巧都值钱。搞懂了计算图和 Autograd 的内存节奏你会发现 PyTorch 里的显存不是玄学而是一本你随时能翻开的账本。
