大模型训练显存瓶颈破解：混合精度与ZeRO优化技术详解！

本文详解大模型训练中的显存优化技术，分析了Adam优化器显存消耗大的原因，介绍了混合精度训练(FP16/BF16)及动态损失缩放技术，重点讲解了ZeRO零冗余优化器如何通过分区优化器状态、梯度和参数来大幅减少显存占用，并提供了PyTorch实现示例，为训练大模型提供实用解决方案。

发菜君

225人浏览 · 2025-10-09 13:38:05

发菜君 · 2025-10-09 13:38:05 发布

在训练大语言模型时，研究人员通常会用 Adam 优化算法。它的好处是比传统的 SGD（随机梯度下降） 更快、更稳定。不过，Adam 除了要保存每个参数的梯度，还要额外存储一阶动量和二阶动量信息，这就让显存（内存）消耗变得非常大。

为了解决显存紧张的问题，大家常用的办法就是 混合精度训练。简单说，就是把不同精度的数字混着用：

FP32（32 位浮点数）：精度最高，最安全，但占内存多。
FP16（16 位浮点数）：只用一半的位数，能节省一半的存储，但数值范围变小，容易出现溢出。
BF16（BFloat16）：和 FP16 一样是 16 位，但指数部分更长，能表示更大的范围，代价是精度比 FP32 差一些。

举个例子，FP32 就像一把长尺子，能量精确地测量大大小小的数值；FP16 的尺子短得多，容易“量不到”；而 BF16 则牺牲了一点精度，换来一个能量更大数值的尺子。

不过，FP16 和 BF16 精度都比 FP32 低，训练过程中可能会出现 梯度消失 或 模型不稳定。为了应对这些问题，研究人员会用一些“辅助工具”：

动态损失缩放：自动调整数值大小，避免溢出或下溢。
混合精度优化器：专门优化计算过程，保证既省内存又不影响训练效果。

这样一来，就能在“省显存”和“保精度”之间找到一个平衡点，让大模型既能跑得动，又能稳定收敛。

混合精度优化（Mixed Precision Optimizer）

以常用的 Adam 优化器 为例，它除了保存模型的参数和梯度之外，还需要维护一些额外的信息，这就是显存消耗巨大的原因。

Adam 优化器都保存了什么？

模型参数：就是网络里的权重，决定模型的计算逻辑。
梯度：反向传播时计算出来的“修正方向”，指导参数该往哪边调整。
一阶动量（Momentum）：可以理解为“梯度的加权平均”。

如果把梯度更新比作推一个小球下山，一阶动量就是“惯性”，它会让小球的速度更加平滑，不会忽快忽慢。
数学上就是累积过去的梯度，起到“记忆”作用。

二阶动量（Variance / RMS）：记录的是“梯度的平方的平均值”。

可以理解为在观察梯度的震动幅度，帮助优化器自动调节学习率。
如果某个方向的梯度变化太大，就会降低步子；变化稳定的方向，则可以走得更快。

这三部分加起来，就构成了 Adam 的核心状态。

为什么显存占用这么大？

假设模型有 Φ 个参数：

参数和梯度用 FP16（16 位浮点数） 存储，分别需要 2Φ 字节。
Adam 里的 参数备份、一阶动量、二阶动量 都用 FP32（32 位浮点数） 存储，各需要 4Φ 字节。
总共算下来，就是 16Φ 字节，其中光是 Adam 的状态就占了 75%。

举个例子：一个 75 亿参数 的大模型，参数如果用 FP16 存储，大约只要 15GB 显存。但是训练时加上 Adam 的状态，实际会吃掉 120GB 显存 —— 这就是为什么显卡总是不够用。

怎么缓解内存压力？

动态损失缩放（Dynamic Loss Scaling）

在反向传播前，把损失（loss）人为放大2^K 倍，这样算出来的梯度也会变大，避免在 FP16 下太小而被当成 0（下溢）。
在更新参数前，再把梯度缩小2^K 倍，还原成正常值。

激活检查点（Activation Checkpointing）

正向传播时不保存所有中间结果，而是只存一部分，等需要反向传播时再临时算回来。
这样能大幅减少激活值占用的内存。

不过，真正的大头还是 Adam 的优化器状态。所以，研究人员一直在探索 更省显存的优化器，例如 ZeRO 优化器、Adafactor 等，让大模型训练更轻便。

零冗余优化**（ZeRO, Zero Redundancy Optimizer）**

在训练超大规模模型时，显存最吃紧的地方往往是 模型状态（参数、梯度、优化器动量等）。为了节省显存，研究人员提出了一种非常巧妙的方法 : 零冗余优化器。

ZeRO 的核心思路

一般情况下，每张显卡（计算设备）都会保存一份完整的模型状态，这就导致大量的重复存储。ZeRO 的做法是：
👉 把模型状态 切分成多个分区，然后分散到不同的设备上。这样整个训练系统里只存在“一份完整的状态”，每张卡只需要保存自己负责的那一部分，就大大减少了显存消耗和通信开销。

ZeRO 的三种分区方式

ZeRO 分为三个阶段（在 DeepSpeed 框架中对应 Zero-1、Zero-2、Zero-3）：

Zero-1：分区优化器状态（Pos）

每张显卡仍然保存完整的 模型参数和梯度。
只有 Adam 的一阶动量、二阶动量等优化器状态被分区。
显存占用：4Φ +12N / Φ 字节（Φ 表示参数量，N 是设备数）。
当 N 很大时，每张卡只需要接近 1/4 的原始内存开销。

Zero-2：分区优化器状态 + 梯度（Pos+g）

每张显卡保存完整的 模型参数，但优化器状态和梯度都分区。
显存占用：2Φ +(2Φ+12Φ) / N 字节。
当 N 很大时，每张卡大约只需 1/8 的原始内存开销。

Zero-3：分区优化器状态 + 梯度 + 参数（Pos+g+p）

所有内容都分区，包括 参数本身。
显存占用：16Φ / N字节。
当 N 足够大时，每张卡的开销几乎趋近于 0。
代价是通信量更高：比常规训练多 1.5 倍的通信量，但相比节省下来的显存，这个代价是可接受的。

总结一下：

Zero-1：解决一部分显存瓶颈，代价小，最稳妥。
Zero-2：进一步节省显存，训练更轻便。
Zero-3：最大化省显存，几乎可以无限扩展模型，但通信开销会增加。

可以说，ZeRO 的出现，让数千亿参数大模型从不可能变成了可能。DeepSpeed 框架就是把 ZeRO 的这套分区理念真正落地实现的代表。

PyTorch中也实现了ZeRO优化方法，可以用ZeroRedundancyOptimizer调用，也可与“torch.nn.parallel.DistributedDataParallel”结合使用，以减少每个计算设备的内存峰值消耗。

使用 ZeroRedundancyOptimizer 的参考代码如下所示：

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.distributed.optim import ZeroRedundancyOptimizer
from torch.nn.parallel import DistributedDataParallel as DDP
def print_peak_memory(prefix, device):
if device == 0:
print(f"{prefix}: {torch.cuda.max_memory_allocated(device) // 1e6}MB")
def example(rank, world_size, use_zero):
torch.manual_seed(0)
torch.cuda.manual_seed(0)
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
# 创建默认进程组
dist.init_process_group("gloo", rank=rank, world_size=world_size)
# 创建本地模型
model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
print_peak_memory("Max memory allocated after creating local model", rank)
# 构建 DDP 模型
ddp_model = DDP(model, device_ids=[rank])
print_peak_memory("Max memory allocated after creating DDP", rank)
# 定义损失函数和优化器
loss_fn = nn.MSELoss()
if use_zero:
optimizer = ZeroRedundancyOptimizer(  # 这里使用了 ZeroRedundancyOptimizer
ddp_model.parameters(),
optimizer_class=torch.optim.Adam,  # 包装了 Adam
lr=0.01
)
else:
optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.01)
# 前向传播
outputs = ddp_model(torch.randn(20, 2000).to(rank))
labels = torch.randn(20, 2000).to(rank)
# 反向传播
loss_fn(outputs, labels).backward()
# 更新参数
print_peak_memory("Max memory allocated before optimizer step()", rank)
optimizer.step()
print_peak_memory("Max memory allocated after optimizer step()", rank)
print(f"params sum is: {sum(model.parameters()).sum()}")
def main():
world_size = 2
print("=== Using ZeroRedundancyOptimizer ===")
mp.spawn(example,
args=(world_size, True),
nprocs=world_size,
join=True)
print("=== Not Using ZeroRedundancyOptimizer ===")
mp.spawn(example,
args=(world_size, False),
nprocs=world_size,
join=True)
if __name__ == "__main__":
main()

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

一直在更新，更多的大模型学习和面试资料已经上传带到CSDN的官方了，有需要的朋友可以扫描下方二维码免费领取【保证100%免费】👇👇

在这里插入图片描述