大模型训练显存瓶颈的常见成因与排查方法

显存瓶颈是大模型训练中最先出现的硬件资源问题，它的直接原因是模型参数、优化器状态、激活值与中间张量在有限的显存空间内叠加超出容量上限。排查显存瓶颈时，应当先看显存分配器的日志，再看Profiler的逐张量内存快照，这两步能定位出九成以上的异常占用。

显存瓶颈的本质是容量和速度的双重天花板

显存瓶颈包含两层含义：一是物理显存容量不足以支撑训练所需的峰值内存，触发Out of Memory错误；二是显存分配和拷贝的速度跟不上计算单元的吞吐需求，导致GPU计算单元空转。这两个因素在实践中同时作用，但容量问题更常见，它直接决定batch size和模型规模的上限。

显存的分配机制决定了碎片化会加剧容量压力。PyTorch的缓存分配器会按块预留显存，训练过程中反复申请和释放张量容易形成碎片。在分布式训练场景下，每个进程的显存预分配还会叠加，进一步收窄可用的连续显存区域。据NVIDIA深度学习框架团队在一次技术分享中的结论，显存碎片可导致高达30%的额外内存消耗。

模型参数和优化器状态是显存占用的固定大头

模型参数和优化器状态在训练开始前就已固定，它们占据的显存不随batch size变化，是显存消耗的底数。一个FP16精度的70B参数模型，仅模型权重就需要140GB显存；若使用Adam优化器，每个参数需要额外存储fp32的动量与方差，加上fp32的权重副本，优化器状态是参数量的12倍字节数。

这个比例可以用一个简单公式表达：训练显存下限 ≈ 参数量 × (权重字节数 + 优化器状态字节数 + 梯度字节数)。以7B模型为例，如果全部用fp32保存，权重28GB，梯度28GB，优化器状态84GB，合计140GB；如果权重和梯度用fp16，优化器状态用fp32，则合计约70GB。这笔固定开销是规划显存预算的第一参考。

激活值显存在反向传播时达到峰值

激活值的显存峰值出现在反向传播阶段，因为框架需要保存每一层的输入张量来计算梯度。激活值显存与batch size、序列长度、隐藏层维度成正比，并且在transformer模型中随层数线性累加。实测中，激活值显存可以占到总显存的40%到60%，在长序列训练场景下甚至超过参数显存。

以一个常见的GPT类模型配置为例，layer=24，hidden=2048，sequence=2048，batch=8，激活值显存估算公式约等于batch×sequence×hidden×(34+5×layers)字节，代入计算约8×2048×2048×(34+5×24)=8×2048×2048×154字节≈4.9GB。当batch增大到32时，激活值显存立刻翻到超过19GB，远超过参数和优化器状态的变化幅度。激活值显存的增长是线性的，但它的基数大，稍微增加batch就能把显存推爆。

用PyTorch Profiler定位显存瓶颈的步骤

PyTorch Profiler提供了内存时间线功能，可以精确记录每个张量的分配、释放和峰值时刻。定位显存瓶颈的标准流程是：先记录内存快照，再按内存占用降序排列张量，最后对照代码行找到源头。整个过程只需要三步。

第一步，启动Profiler并记录内存事件。代码示例如下：

<pre><code>from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof: train_one_step(model, batch) print(prof.key_averages().table(sort_by=