标签:显存优化
-
ZeRO:把优化器状态也切了,显存占用的递进削减方案
朴素数据并行每卡都存一份完整的权重、梯度和 Adam 状态,对 7B 模型这是近 200 GB 的浪费。ZeRO 的思想朴素:既然梯度反正要求平均,分片存储、用时再聚合。三个优化阶段分别对应分片优化器状态、梯度、权重,把单卡显存从 O(Ψ) 一路压到 O(Ψ/N),代价是通信量的有界上涨。这篇算清每阶段省多少、多花多少通信,以及 PyTorch FSDP 对应的是哪一档。
-
梯度累积与有效 batch size:显存换批量的等价性边界
显存装不下目标 batch size 时,梯度累积把 k 个 micro-batch 的梯度加起来再走一步优化器,数学上近似等价于大 batch 训练。这篇推导这个等价的条件,拆解归一化方式、梯度裁剪、BN 统计量三个最常踩的坑,并说明它和多卡数据并行的组合关系。
-
梯度检查点:用 33% 的额外计算换显存平方级的节省
反向传播需要前向的所有激活,显存随层数线性增长——这是长序列训练的第一堵墙。梯度检查点只存每个分段的边界激活,反向时重算段内前向,把激活显存从 O(L) 压到 O(sqrt L),代价约 33% 的额外 FLOPs。这篇推导最优分段数的公式,讲清选择性重计算(Megatron 的做法)为什么比全量重算更划算,以及 torch.utils.checkpoint 的实现细节。