标签:ZeRO
-
ZeRO:把优化器状态也切了,显存占用的递进削减方案
朴素数据并行每卡都存一份完整的权重、梯度和 Adam 状态,对 7B 模型这是近 200 GB 的浪费。ZeRO 的思想朴素:既然梯度反正要求平均,分片存储、用时再聚合。三个优化阶段分别对应分片优化器状态、梯度、权重,把单卡显存从 O(Ψ) 一路压到 O(Ψ/N),代价是通信量的有界上涨。这篇算清每阶段省多少、多花多少通信,以及 PyTorch FSDP 对应的是哪一档。
朴素数据并行每卡都存一份完整的权重、梯度和 Adam 状态,对 7B 模型这是近 200 GB 的浪费。ZeRO 的思想朴素:既然梯度反正要求平均,分片存储、用时再聚合。三个优化阶段分别对应分片优化器状态、梯度、权重,把单卡显存从 O(Ψ) 一路压到 O(Ψ/N),代价是通信量的有界上涨。这篇算清每阶段省多少、多花多少通信,以及 PyTorch FSDP 对应的是哪一档。