大模型训练显存优化：原理、技术与实践

1. 大模型显存消耗的本质与构成

在深度学习模型训练过程中，显存消耗主要来自三个核心部分：模型参数本身、计算图（DAG）和临时栈变量。理解这三者的构成和比例关系，是进行显存优化的基础。

模型参数部分包含三个关键组成部分：

• 权重参数（weight）：模型的可训练参数，如线性层的权重矩阵
• 梯度（gradient）：反向传播时计算的参数梯度
• 优化器状态（opt states）：包括动量（m）、二阶动量（v）和主权重（master weight）

以bf16混合精度训练为例，这三者的典型比例为1:1:6。这意味着对于一个72B参数的模型，仅优化器状态就需要约432GB的显存空间，这解释了为什么大模型训练需要如此多的显存资源。

计算图（DAG）是显存消耗的第二大来源。在前向传播过程中，需要保存中间激活值以实现自动求导。这些激活值会随着batch size和序列长度的增加而线性增长，对于深层网络来说可能占据数百GB的显存。

临时栈变量则包括：

• 所有操作的输入输出张量
• 计算过程中的中间缓冲区
• 短期使用后会被Python垃圾回收机制释放的内存

提示：在实际训练中，临时变量的显存管理往往被忽视，但它们可能导致严重的显存碎片问题，特别是在处理大张量时。

2. 显存优化的核心价值与收益

显存优化不仅仅是让大模型能在有限硬件上运行的技术手段，它还能带来多方面的实际收益：

资源利用率提升：通过优化，可以使用更少的GPU卡完成原本需要更多硬件资源的训练任务。例如作者实现的72B模型在8卡80G配置下的训练，相比传统方法节省了50%以上的硬件投入。

训练策略灵活性增强：

• 可以灵活调整Tensor Parallelism(TP)/Pipeline Parallelism(PP)的规模
• 减少TP规模可以降低AllReduce通信开销
• 优化PP配置可以减少流水线气泡（bubble）时间
• 避免使用Checkpointing(CP)策略，防止显存需求翻倍

训练效率提升：

• 减少recompute层数
• 增大micro batch size
• 在计算核心未被充分利用时提高MFU（Model FLOPs Utilization）
• 保持访存和通信开销不变的情况下提升吞吐量

工程实践价值：