DeepSpeed：PyTorch优化库，使模型分布式训练能高效使用内存和更快速

时间：2024-11-14 11:41:40

简介

DeepSpeed是一个PyTorch优化库，它能让大模型的分布式训练更加节省内存和更快速。它的核心是ZeRO,ZeRO让大规模的训练大模型成为可能。ZeRO有以下几个工作状态：

  • ZeRO-1：通过多个GPU进行 优化器状态 分布计算
  • ZeRO-2：通过多个GPU进行 梯度 分布计算
  • ZeRO-3：通过多个GPU进行 参数 分别计算

在GPU受限的环境下，ZeRO可以将优化器内存和计算从GPU卸载到CPU从而利用单个GPU训练大模型。DeepSpeed 已经集成到Transformers Trainer里面完成ZeRO所有的状态和卸载。我们需要做的只是提供一个配置文件或者使用一个已经提供好的文件模版。针对推理，Transformers支持ZeRO-3和卸载，从而允许加载超大模型。

本文将会告诉我们如何通过DeepSpeed进行训练，哪些特征我们可以启用，如何去配置不同的ZeRO状态、卸载、推理，并在不使用Trainer的情况下使用DeepSpeed。

安装

DeepSpeed可以通过PyPI或者Transformers 进行安装，具体如下：

pip install deepspeed
pip install transformers[deepspeed]

如果在安装DeepSpeed的时候遇到了困难，可以查看指导手册： DeepSpeed CUDA installation。当然DeepSpeed也可以通过PyPI包进行安装，这也是我们极力推荐的安装方式，通过源码 install it from source 进行安装能最好的适配我们的硬件和支持某些特性，比如1-bit Adam,这个通过PyPI安装是不可用的。

内存要求

在我们开始之前，一个好的策略是，检测一下我们是否有充足的GPU和CPU内存够加载我们的模型。DeepSpeed提供了一个工具来评估需要的CPU/GPU内存。例如，评估单个GPU上面运行bigscience/T0_3B模型需要的内存：

$ python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("bigscience/T0_3B"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1)'
[...]
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 1 GPU per node.
SW: Model with 2783M total params, 65M largest layer params.
  per CPU  |  per GPU |   Options
   70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
   70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
   62.23GB |   5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=1
   62.23GB |   5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=0
    0.37GB |  46.91GB | offload_param=none, offload_optimizer=none, zero_init=1
   15.56GB |  46.91GB | offload_param=none, offload_optimizer=none, zero_init=0

num_gpus_per_node=1：一台机器上面有几个GPU卡

num_nodes=1：当前集群有几台机器

上面的评估结果显示；要么我们选择单个80G GPU不需要CPU资源，要么选择8GP GPU和 0~60GB CPU的装载。（这里的内存仅仅是参数、优化器状态和梯度的消耗，我们还需要考虑CUDA内核和活动消耗的内存）我们需要在成本和速度上做一下取舍，如果我们使用一个比较小的GPU就需要占用比较多的时间来训练我们的模型。

如果我们有充足的GPU内存来确保我们不用CPU/NVMe装载，那一切都是最快的。

选择ZeRO状态

当我们安装好DeepSpeed，并对使用我们自己的内存需求有一个大致的了解后，下一步就是选择要使用哪种ZeRO状态。为了更快和更有效率的使用内存：

Fastest Memory efficient
ZeRO-1 ZeRO-3 + offload
ZeRO-2 ZeRO-3
ZeRO-2 + offload ZeRO-2 + offload
ZeRO-3 ZeRO-2
ZeRO-3 + offload ZeRO-1

选择一个最适合我们的组合，最开始我们使用最快的组合直到我们出现OOM，然后选择下一个状态，它将会慢一点但是更节省内存。从节省内存或者最快速度可以随意的选择哪一种，最终选择一个最适合的速度和内存使用的平衡点。

一个通常的步骤供我们借鉴（从 batch size 等于 1 开始）:

  1. 开启梯度 checkpointing
  2. 尝试ZeRO-2
  3. 尝试ZeRO-2 和 卸载优化器
  4. 尝试ZeRO-3
  5. 尝试ZeRO-3和 卸载参数到CPU
  6. 尝试ZeRO-3 和 卸载 参数和优化器到CPU
  7. 尝试降低各种默认值，比如：如果我们使用了generate()方法，则收窄搜索窗口
  8. 尝试在完全精度权重中混合半精度（在老的GPU架构上用FP16和在Ampere上使用bf16）
  9. 使用更多的硬件，或者将 参数和优化器卸载到NVMe
  10. 一旦我们不在出现OOM，调节有效吞吐量，然后增加 batch size 直到我们将GPU效率最大化
  11. 最后，尝试优化我们训练的启动参数：禁用一些卸载特性或者使用更快的ZeRO阶段，增加/减小 batch size 来发现速度和内容使用的最佳平衡点。

下一篇文件将会介绍DeepSpeed配置文件

相关文章