大模型训练如何节省显存 - DeepSpeed ZeRO

参考

https://www.deepspeed.ai/tutorials/zero/#zero-overview

https://mmengine.readthedocs.io/zh-cn/latest/common_usage/large_model_training.html

介绍

训练大模型需要巨大的资源，单卡显存已经越来越难以满足存放整个模型，因此诞生了大模型训练技术，典型的如 DeepSpeed ZeRO 和 FairScale 的完全分片数据并行（Fully Sharded Data Parallel, FSDP）技术，其允许在数据并行进程之间分片模型的参数、梯度和优化器状态，并同时仍然保持数据并行的简单性。

DeepSpeed的Zero Redundancy Optimizer（ZeRO）通过在多个GPU之间分割模型的权重、梯度和优化器状态，显著减少了每个GPU所需的内存量。

这使得训练非常大的模型成为可能，而这些模型在没有ZeRO的情况下可能会因为内存限制而无法训练。

例子：训练一个1.5B参数的GPT-2模型

假设我们有一个具有1.5亿（1.5B）参数的GPT-2模型，我们想要在一个由8个NVIDIA Tesla V100-SXM3 GPU组成的集群上训练它。每个GPU有32GB的RAM。

没有ZeRO的情况：

• 我们尝试在每个GPU上使用默认的数据并行设置来训练模型。
• 由于模型的参数和优化器状态（例如Adam优化器的一阶和二阶矩估计）需要存储在每个GPU上，这会占用大量的内存。
• 结果是，每个GPU的内存需求超过了其可