ZeRO-1/2/3 显存切分

ZeRO(Zero Redundancy Optimizer, DeepSpeed)通过切掉“冗余保存”把大模型训练塞进有限的 GPU 显存。设参数总数为 $\Psi$(FP16 每元素 2 字节),用 Adam 混合精度训练需同时保存 FP16 参数 $2\Psi$、反向产生的 FP16 梯度 $2\Psi$,以及 FP32 优化器状态——主权重副本 $4\Psi$ + 一阶动量 $4\Psi$ + 二阶动量 $4\Psi$ 共 $12\Psi$,静态显存合计 $16\Psi = 2\Psi + 2\Psi + 12\Psi$。70B 模型: $16 \times 70$ GB $= 1120$ GB,单张 H100(80 GB)装不下,朴素 DDP 在 $N_d$ 卡上复制 $N_d$ 份完整状态、更装不下。ZeRO 分三档切分(设 $N_d$ 卡,每卡显存): ① ZeRO-1 只切优化器状态: $2\Psi + 2\Psi + \tfrac{12\Psi}{N_d}$;② ZeRO-2 再切梯度: $2\Psi + \tfrac{2\Psi + 12\Psi}{N_d}$;③ ZeRO-3 连参数也切: $\tfrac{16\Psi}{N_d}$——70B 需 $N_d = 16$ 张 H100(每卡 $1120/16 = 70$ GB)才放得下。还不够时 Offload: 优化器状态与梯度移到 CPU 内存,GPU 只保留本层 FP16 参数分片 $2\Psi/N_d = 140/8 = 17.5$ GB(8 卡),必要时按页换到 NVMe(ZeRO-Infinity,可扩展至万亿参数)。

大模型预训练/微调(7B–175B)显存预算的面试计算题;高频追问“手算 70B 训练显存”“ZeRO-2 相比 ZeRO-1 省什么”“Offload 换来什么代价”。

朴素 DDP 在 $N_d$ 卡上冗余保存 $N_d$ 份完整状态,显存不随卡数下降;ZeRO 把冗余从 $O(N_d)$ 倍压到 $O(1)$ 倍,每卡显存降一个数量级(70B: 1120 GB → 每卡 70 GB),通信量只比 DDP 多约 1.5 倍(ZeRO-3 约 3Φ vs DDP 的 2Φ),配合 CPU/NVMe Offload 甚至能在单节点训练 175B 级模型。

手算 70B 模型(Adam + FP16/FP32 混合精度)训练的单卡静态显存: 写出 $16\Psi = 2\Psi + 2\Psi + 12\Psi$,逐项说明(FP16 参数 / FP16 梯度 / FP32 主权重 + 一阶 + 二阶动量),得出 1120 GB,并解释为何朴素 DDP 每卡仍是 1120 GB(复制 N 份冗余)。

写出 ZeRO-1/2/3 三档的每卡显存公式 $2\Psi + 2\Psi + \tfrac{12\Psi}{N_d}$、$2\Psi + \tfrac{2\Psi + 12\Psi}{N_d}$、$\tfrac{16\Psi}{N_d}$,并说明每一档切了什么、还剩什么未切。

ZeRO-2 相比 ZeRO-1 省了什么?省多少(梯度从每卡 $2\Psi$ 降到 $2\Psi/N_d$)?梯度为什么能被切——反向后立即 Reduce-Scatter 聚合分片,而不是全量 All-Reduce。

ZeRO-3 的通信量为什么约是 DDP 的 1.5 倍(前向 All-Gather 参数 Φ + 反向 Reduce-Scatter Φ + 反向重算时再 All-Gather 参数 Φ ≈ 3Φ,vs DDP 的 2Φ),显存收益 $16\Psi \to 16\Psi/N_d$ 是否划算。

ZeRO-Offload 把什么移到 CPU、代价是什么: 优化器状态与 Adam 计算下放 CPU,GPU 只留本层参数分片 $2\Psi/N_d$;换来显存大降,代价是 PCIe 带宽受限、步时间变长;ZeRO-Infinity 如何再扩展到 NVMe。