1 显存优化要回答什么
训练图像分类模型时，程序可能顺利读完数据，却在第 1 次反向传播时报出显存不足，即 OOM。也可能前面训练了 200 个批次都正常，碰到较长文本或执行验证时才失败。显存优化要回答的不是“哪个开关最省”，而是：什么数据在什么时刻占据显存，哪些必须保留，哪些能少存、晚算或移到别处，以及这些改变会怎样影响训练速度和模型质量。
本课以公开的 CIFAR-10 图像分类任务为主线。输入是 32×32 的彩色图像，输出是 10 个类别的预测概率。将官方训练集的 50000 张图像固定划分为 45000 张训练图像与 5000 张验证图像，官方 10000 张测试图像保留到最终评估。训练集可以随机裁剪、翻转，验证集与测试集采用固定预处理。划分索引和随机种子必须保存，否则每次换显存策略又换验证样本，就无法公平比较。
设训练设备标称显存为 8 GiB，当前模型使用 Adam 优化器。我们希望在这块设备上稳定训练，记录每轮耗时、训练峰值显存、验证集宏平均 F1，以及各类别的精确率和召回率。宏平均 F1 先分别计算各类别的 F1，再取平均，使每个类别都得到同等关注。若某个类别经常被错分，只报告总体准确率可能掩盖问题。显存策略通过改变批大小或数值精度，也可能间接改变这些指标，因此“程序能跑”只是起点。
同样的分析适用于表格客户流失预测和短文本分类。前者常见特征包括任期、月费用和历史服务使用情况；预测时点之后才出现的解约原因不能作为输入，否则会发生标签泄漏。后者的显存压力常随批次中最长文本增长。类别不平衡时，需要同时观察少数类召回率、F1 或适合任务的 AUC；部署时期的数据分布发生变化时，也不能把指标下降全部归咎于显存策略。硬件问题和评估问题要分别建立证据。
先明确 3 个约束：允许使用的设备、可以接受的训练时间、必须达到的质量标准。接着找到峰值发生在哪个阶段，再选择策略。若峰值来自激活，减少优化器状态可能帮助有限；若模型还没执行前向就放不下，只缩小输入批次也不能解决根因。显存优化的核心是让措施对应主要占用，并用可复现的实验验证。
2 激活、参数与优化器状态
把显存理解为训练的工作台。参数是模型需要学习的权重，前向传播会反复读取它们；梯度记录损失对参数的变化方向，供参数更新使用；优化器状态是更新规则保存的历史信息，例如 Adam 为每个参数维护的 2 组统计量。它们通常随模型规模增长，与当前批次包含多少样本没有直接的线性关系。即使批大小降到 1，这部分占用也不会自动消失。
激活是前向计算产生的中间结果。为了在反向传播中计算梯度，框架通常要保留部分输入和中间张量。批大小、图像分辨率、网络层数、隐藏维度，以及具体算子的实现，都会影响激活占用。激活不是“最后的模型输出”的同义词：最后只有 10 个类别概率，内部仍可能保存大量特征图。因此，仅根据输出张量大小估算训练显存，往往会低估很多。
训练显存还包括输入、损失计算的临时张量、算子工作区和运行时开销。某些临时张量只存在很短时间，却可能恰好与尚未释放的激活重叠，形成峰值。评估通常不需要保存训练用的反向图，但必须明确关闭梯度记录；只调用模型的评估模式，并不等于关闭自动求导。评估模式主要改变某些层的行为，是否记录梯度由另外的机制控制。
可以先写近似关系：峰值显存 ≈ 参数与梯度及优化器状态 + 峰值存活激活 + 临时空间与运行时开销。这个关系帮助分类，不能把各阶段各自测得的最大值机械相加，因为它们未必同时存在。参数更新时优化器可能创建临时缓冲区，反向时激活逐步释放；需要沿时间观察，才能知道真实最高点在哪里。
还要区分“已分配”和“已保留”。在常见缓存分配器中，已分配表示当前活跃张量使用的空间，已保留表示框架向设备取得并管理的空间，后者可能包含可复用空闲块。设备监控工具看到的进程占用还可能包含框架统计之外的开销。已保留大于已分配并不自动证明内存泄漏，也不自动证明碎片化。判断必须结合曲线、分配失败位置和可复现实验。
动画沿着前向、反向和更新移动时间指针。观察重点是每段空间什么时候出现、什么时候可以复用，而不是记住固定比例。不同模型的占用结构差异很大；“激活总占多少百分比”没有适用于所有任务的答案。
3 重计算、分片与卸载
当激活是主要占用时，可以考虑激活重计算，也叫激活检查点。普通训练保存较多中间结果；重计算只保存选定边界，在反向需要某些中间值时，再执行相应的前向计算。直观地说，是用额外计算换取较少的长期保存空间。模型参数仍然要存在，梯度和优化器状态也不会因为开启重计算就全部缩小。
检查点放得越密不一定越省，放得越少也不一定越快。边界张量本身需要保存，区域内部重算需要时间，某些区域的计算成本远大于它节省的空间。入门实践可以先对占用较大的重复模块分组试验，不必从每个算子的内部实现入手。包含随机行为或会修改内部状态的模块，应使用框架支持的检查点接口，并检查相应限制，确认重算保持了所需的训练语义。
分片针对的是另一类重复占用。普通数据并行常让每个设备各存完整参数、梯度和优化器状态；分片方案把其中部分状态分散到多个设备。需要使用时再通过通信取得相应数据。分片能让单个设备承担更少的长期状态，但会引入数据交换和临时聚合的空间。不能把模型总字节数除以设备数，就断言每个设备的实际峰值；通信缓冲区、激活和聚合时机仍然存在。
卸载则把暂时不用的数据移到主机内存等位置，需要时再取回。它扩大了可使用的存储范围，但传输链路有带宽和延迟限制，主机内存也有容量上限。显存下降而每轮训练明显变慢，可能说明程序正在等待数据搬运。卸载优化器状态与卸载激活的时序不同，不能把它们看作相同开关。具体接口与异步调度留到实际框架专题，本课先掌握“放在哪里、何时取回、谁在等待”的判断方式。
这些策略可以组合，但先后顺序由测量决定。单个设备上激活偏大时，先减小微批次、检查数值精度或试验重计算通常更直接；多设备训练的状态重复明显时，再评估分片。若单个设备完全装不下模型状态，就不能假设梯度累积一定救得回来。每加 1 种机制，都要重新测峰值、速度和质量，避免多个变化叠加后无法解释结果。
4 动手：估算一次模型显存账
现在为实验模型做预算。假设它有 20000000 个可训练参数，参数和梯度都使用 FP32，每个数占 4 字节；Adam 的 2 组状态也使用 FP32。先忽略实现额外开销，则参数占 80000000 字节，梯度占 80000000 字节，优化器状态占 160000000 字节，合计 320000000 字节，约 305 MiB。这里 1 MiB 为 1048576 字节，1 GiB 为 1024 MiB。不要把十进制 GB 和 GiB 混在同一张预算表中。
这个数字只是长期状态的简化估算。部分优化器采用不同实现，混合精度训练也可能保留 FP32 参数副本或不同精度的状态。把所有张量都按 2 字节计算，会低估实际占用。正确做法是先列出本次配置中每类张量的精度和份数，再估算；如果拿不到细节，就将这部分标注为待测，而不是给出貌似精确的总数。
在固定输入尺寸下，假设实测微批大小为 16、32、64 时，峰值已分配显存分别约为 1.8、3.0、5.4 GiB。这组数字是教学示例，不代表任何特定模型的性能。相邻观测表明，每增加 1 个样本，峰值大约增加 0.075 GiB，固定部分约为 0.6 GiB。固定部分大于前面的状态估算并不矛盾，因为还有临时空间和其他常驻数据。这个局部线性关系可以帮助挑选下一次实验，却不能保证跨分辨率、跨算子实现仍然成立。
若设备为 8 GiB，不应直接把模型预算填到 8 GiB。显示占用、其他进程、验证阶段和偶发的大输入都可能需要空间。可以先把训练峰值预算设为 6.5 GiB，再根据多次完整运行调整。按当前近似，微批大小 80 的峰值约为 6.6 GiB，已超过这条实验预算线；先测试 64 更稳妥。预算线是本次实验选择，不是所有设备通用的固定比例。
实际测量应从干净进程开始，固定模型、数据划分、输入形状和精度，先做若干预热迭代，使延迟创建的优化器状态和算子准备工作出现。随后重置峰值统计，测量完整训练迭代，包括前向、损失、反向和参数更新。设备操作常异步执行，测耗时时要在测量边界使用框架支持的同步方式；不必在每个算子后同步，否则会改变原本的执行节奏。
再单独测验证阶段，并记录已分配峰值、已保留峰值、设备总占用、每步耗时和样本吞吐量。日志里写明软件版本、设备、微批大小、输入形状、优化器和精度。先做短程容量试验，再跑完整训练；短程试验通过只能说明当前样本和阶段暂时可运行，不能证明后续长文本、验证或保存检查点时一定安全。
5 速度与显存取舍
缩小微批次通常能降低激活占用，但可能降低设备利用率。若希望保持每次参数更新使用的样本数，可以采用梯度累积：连续处理多个微批次，累积梯度后再更新参数。在单个设备上，微批大小 16、累积 4 次，对应有效批大小 64。多个设备采用相同累积方式时，还需要乘以参与数据并行的设备数。有效批大小描述更新覆盖的样本数，峰值激活主要受实际同时处理的微批次影响。
实现累积时要检查损失缩放。若各微批次样本数相同，常见做法是把每次平均损失除以累积次数，再反向传播；最后不足完整组的微批次应按实际样本数处理。梯度应在适当的更新边界清空。把梯度每次都清空，会失去累积效果；完全忘记清空，则会跨更新不断叠加。学习率调度也要明确按微批次推进还是按参数更新推进，避免悄悄改变训练计划。
即使有效批大小相同，累积训练也不保证逐位等价于直接大批次训练。归一化层看到的微批次统计、随机层和浮点运算顺序都可能不同。对图像任务，应观察各类别召回率与宏平均 F1，而不是只看损失是否下降。若小微批次让含批次统计的模型表现不稳定，就需要单独评估相关层的处理，不能只增加累积次数后认定问题已经解决。
混合精度可减少部分张量的存储并可能提高速度，但不是把所有占用统一减半。需要设备与算子支持合适的数据格式，并监控非有限损失、梯度异常和质量变化。某些数值格式和训练配置需要梯度缩放。以当前框架推荐的自动混合精度接口为起点，比手工把每个张量都转换成低精度更容易保持正确性。验证时也要记录所用精度，确保比较条件清楚。
滑块展示一个教学模型：微批次越大，激活峰值上升；重计算覆盖越广，长期保存的激活下降，但计算时间上升。它用于理解方向，不替代真实基准。选配置时应比较在相同训练任务和更新计划下的每轮时间、峰值显存和验证指标。吞吐量最高的配置如果持续 OOM，就没有实用价值；节省最多显存的配置如果使训练时间超过预算，也未必适合当前方案。
6 常见坑：盲目重计算反而更慢
假设微批大小 32 时已经稳定运行，峰值为 3.0 GiB，每轮训练用时 100 秒。把所有模块都加入重计算后，峰值降至 2.2 GiB，每轮却变成 138 秒。如果省下的空间没有用来容纳更大的输入、模型或微批次，训练目标也没有因此改善，那么这项优化只降低了当前并不紧张的资源使用，却增加了时间成本。不能只展示显存下降的百分比，就宣布整体更优。
重计算还可能因为切分过细而收益很小：每段仍要保存边界，额外调用和重复执行却不断增加。切分过粗也可能让某个重算区域的临时峰值偏高。先选激活明显的大模块进行有限试验，再比较收益，通常比遍历所有开关更容易解释。不要在不了解模块行为时绕过框架检查，尤其要注意随机状态、原地修改和依赖外部状态的计算。
另一种误判来自计时不完整。只测前向传播，会看不到反向阶段增加的重算；只测训练开头，则可能混入首次初始化成本。对比时应使用相同设备、相同输入序列与相同更新次数，预热后测完整迭代，报告多次观测的范围。若配置改动导致有效批大小变化，训练到相同轮数也未必具有相同更新次数，需要在记录中说明。
还要防止把保留计算图误认为“模型太大”。例如每步把带有计算图的损失张量直接加入长期列表，或把全部训练输出保存在设备上，可能导致活跃张量持续增长。日志通常只需要脱离计算图的标量，预测结果可以按批转到主机端或累计统计量。使用没有梯度记录的验证过程，也能避免为不需要反向传播的计算保留图。
训练质量必须独立验收。若改完配置后验证 F1 突然显著提高，应先检查验证集是否意外参与训练、数据增强是否改变、标签是否混入输入。对客户流失任务，随机划分还可能掩盖同一客户重复记录或时间泄漏。对短文本任务，不同划分之间的重复文本也会让结果虚高。显存下降是真实的，并不能自动证明指标提升同样可信。
7 案例：碎片化导致 OOM
把主线扩展到公开的 AG News 短文本分类数据。任务是根据标题与描述判断 4 个新闻类别，特征是分词后的词元序列及其有效位置掩码。将官方训练集固定划分为训练部分与验证部分，保留官方测试集。设置最大长度，记录截断比例，采用动态填充，让同批文本补到当前批次最长长度。这样通常能减少无效计算，但不同批次的张量形状会变化。
一次运行中，短文本批次与长文本批次交替出现，随后在验证阶段 OOM。日志显示已分配约 5.1 GiB、已保留约 7.2 GiB，设备容量为 8 GiB。开发者看到已保留与已分配的差额，怀疑碎片化。这个怀疑合理，但数字本身还不够：验证批次是否更大、某条文本是否超过预期长度、其他进程是否占用设备，都可能造成相似现象。不同分配器的行为也有差异。
碎片化可以理解为可复用空间分散在不合适的块中，当前请求未必能利用它们。分配器可能回收缓存或采用不同策略，因此不能简单地说“剩余空间分散就必然失败”。需要查看当前运行环境支持的内存摘要或快照，把失败请求大小、活跃块、不可复用空间与输入形状联系起来。本例在多次复现后发现：相同最大长度的固定形状序列能完成，长度频繁变化的序列则更容易失败，快照也支持分配块利用不佳的判断。
处理时先按长度分桶，使相近长度的样本组成批次，再限制每批词元预算，避免仅按文本条数控制负载。分桶仍要保留适当随机性，防止训练顺序长期按长度固定。最大长度只能根据训练数据分析与验证结果确定，不能为了降低显存随意截掉决定类别的关键信息。验证集和测试集采用明确、一致的长度规则，并报告截断率及长文本子集的表现。
修改后，完整训练与验证均能完成，峰值已保留空间降低，宏平均 F1 和各类别召回率处于基线波动范围内。这才支持将新批处理规则纳入方案。清理缓存有时能释放空闲缓存块，却不能释放仍被活跃张量引用的空间，也不会降低模型本身的必要占用；不应在每步无条件调用。只有测量仍指向分配器行为时，才考虑对应版本支持的配置，并保留单变量对照。
8 把显存策略写进训练方案
训练方案要从任务和基线写起：数据集版本、划分索引、输入处理、模型结构、优化器、精度和评价指标。CIFAR-10 的基线可以是固定随机种子下的完整训练，记录验证宏平均 F1 与各类别召回率；最终测试集只用于冻结方案后的评估。比较优化配置时使用相同划分，必要时重复多个种子，以免把训练随机波动误判为策略影响。
接着写容量预算和实验顺序。例如先测微批大小 32 的完整迭代，确认峰值主要来自激活，再测试混合精度和微批大小 64；若仍不满足目标，选择部分模块做重计算。每次只改变能够单独解释的因素，记录峰值、每轮耗时、有效批大小和质量变化。分片与卸载则附加记录设备数、主机内存峰值与数据传输等待情况，避免只把设备显存转移成另一个未计量的瓶颈。
验收应同时覆盖稳定性、时间与质量。稳定性要求完成训练、验证及保存等必要阶段；速度要求在约定环境和预算内完成；质量要求主要指标及关键类别表现不低于事先设定的容忍范围。容忍范围应来自业务目标和基线波动，不应在看到结果后临时修改。类别不平衡的流失预测还要说明阈值如何在验证集上选择，防止只用总体准确率掩盖少数类召回下降。
最后写失败时的回退路径。输入出现分布变化时，记录长度、分辨率和批次构成，优先回到已验证的小微批次配置；若活跃占用持续增长，先检查计算图和长期引用；若怀疑碎片化，则收集可复现输入与分配器证据。配置文件应保存有效批大小的计算方式和更新节奏。方案的价值在于别人能够按相同步骤复现结果，也能知道何时应该停止沿用当前结论。
9 从显存优化走向集合通信（A43）
显存优化走到多个设备时，原本只需要考虑“保存多少”的问题，会增加“由谁保存、什么时候交换”的问题。数据并行通常让各设备处理不同样本，然后汇总梯度，使参数更新保持一致；分片训练则可能在计算前获取参数片段，在计算后分散梯度或状态。节省的空间来自减少重复保存，代价之一是设备间的数据交换。
例如 2 个设备各处理微批大小 16，并各自累积 2 次，数据并行下的有效批大小为 64。但各设备并不是各自训练到结束再随意平均模型；具体同步边界决定了它们如何共同完成更新。通信发生得太频繁，可能增加等待；推迟通信又可能保留更多数据，或要求框架提供特定的累积机制。因此，单个设备上最快的微批次配置，迁移到多个设备后未必仍然最快。
下一课 A43 将介绍集合通信：多个参与者按照约定共同交换或归约数据，例如把各设备的梯度相加，或者把分散的片段收集起来。本课只需带走 1 条连接：分片减少本地长期占用，却可能增加通信和临时空间。评估多设备方案时，要把显存峰值、计算时间和通信等待放在同一条时间线上，才能解释设备为什么空闲，也才能判断增加设备究竟解决了什么问题。
放进方案里
在训练方案中增加“显存预算与回退”段落，先写清任务：使用 CIFAR-10 的 45000 张训练图像与 5000 张验证图像，官方测试集在方案冻结后使用；输入保持 32×32，训练与评估的预处理分别固定。记录模型版本、优化器、数据精度和设备容量。以验证宏平均 F1 为主要质量指标，同时检查各类别召回率，避免显存调整后只有总体指标看似正常。
容量记录应包含微批大小、累积次数、有效批大小、已分配峰值和已保留峰值。先在干净进程中预热，再测完整训练迭代，随后覆盖验证和保存阶段。设备为 8 GiB 时，可先设 6.5 GiB 的训练峰值预算，并根据完整运行结果调整。保留预算依据和实际观测，不能把预留空间当作跨设备通用常数。
策略按证据推进：激活占用明显时，依次比较微批次调整、混合精度和局部重计算；状态放不下时，再评估分片或卸载。每项实验记录每轮耗时、主机内存需求和质量变化。候选配置必须完成全部必要阶段，在约定时间内训练，并满足提前设定的质量容忍范围，才可替换基线。
验收附件保留峰值所在迭代的输入形状、显存曲线和完整运行日志，并注明是否包含首次优化器更新。容量结论仅适用于已测的分辨率、微批大小与精度组合；更换优化器或增大输入后，必须重新测量，不能沿用原验收结果。
回退条件写成可操作规则：遇到 OOM，保留失败批次的形状和阶段，先恢复已验证的小微批次；若占用逐步上升，检查计算图与长期引用；若形状变化触发失败，结合快照判断是否需要长度分桶或词元预算。任何改动都保存配置、划分索引和随机种子。线上或后续数据变长时，重新验证容量和长样本召回率，不把旧预算直接视为保证。
10 小结：下一步集合通信（A43）
显存账需要同时看长期状态、激活和临时空间，还要沿着训练时间线寻找重叠峰值。减少微批次、重计算、分片和卸载解决的占用来源不同；合适的配置应同时满足容量、速度和质量要求。显存读数下降只是观测结果，完整任务能够稳定完成才具有实际意义。
自测：为什么梯度累积不能解决参数状态本身放不下的问题？为什么已保留远大于已分配还不足以证明碎片化？如果重计算省下 1 GiB，却使每轮训练增加 30 秒，你还需要哪些约束才能作出选择？尝试用本课的测量记录回答，再进入 A43，理解多个设备如何共同交换梯度与参数片段。