大模型训练如何节省显存 - 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的内存需求超过了其可用内存（32GB），导致内存不足（OOM）错误。

使用ZeRO的情况：

• 我们通过修改DeepSpeed配置文件来启用ZeRO优化。
• 在配置文件中，我们设置zero_optimization键，将stage设置为1，并设置reduce_bucket_size为500MB。
• ZeRO将优化器状态分割到不同的GPU上，这样每个GPU只需要存储一部分优化器状态。
• 通过这种方式，每个GPU的内存需求从18GB降低到2.25GB，使得模型可以在8个GPU上顺利训练，而不会遇到内存不足的问题。

ZeRO的不同阶段

ZeRO分为几个阶段，每个阶段都提供了额外的优化：

• ZeRO-1：优化器状态被分割。
• ZeRO-2：模型的3梯度被分割。
• ZeRO-3：模型参数被分割，并且可以自动收集和分割。

此外，ZeRO-3还包括了Infinity offload引擎，它可以将模型状态卸载到CPU和NVMe内存中，从而实现巨大的内存节省。

总结

ZeRO技术使得训练大型模型变得更加可行，特别是在内存受限的情况下。通过减少每个GPU需要存储的数据量，ZeRO允许我们在现有的硬件上训练更大的模型，而无需进行代码修改。这为研究人员和开发者提供了更大的灵活性和可能性，以探索和训练先进的深度学习模型。