ZeRO-Offload

ZeRO-3 Offload 包含我们新发布的 ZeRO-Infinity 中的一部分功能。阅读我们的 ZeRO-Infinity 博客 以了解更多信息！

我们建议您在完成本教程之前，阅读有关 入门 和 ZeRO 的教程。

ZeRO-Offload 是一种 ZeRO 优化，它将优化器内存和计算从 GPU 卸载到主机 CPU。ZeRO-Offload 使得能够在单个 GPU 上高效训练参数高达 130 亿的大型模型。在本教程中，我们将使用 ZeRO-Offload 在 DeepSpeed 中训练一个 100 亿参数的 GPT-2 模型。此外，在 DeepSpeed 模型中使用 ZeRO-Offload 非常快速和简单，因为您只需要更改 DeepSpeed 配置 json 中的一些配置。无需更改代码。

ZeRO-Offload 概述

对于大型模型训练，诸如 Adam 之类的优化器可能会消耗大量的 GPU 计算和内存。ZeRO-Offload 通过利用主机 CPU 上的计算和内存资源来执行优化器，从而减少了此类模型的 GPU 计算和内存需求。此外，为了防止优化器成为瓶颈，ZeRO-Offload 使用 DeepSpeed 高度优化的 CPU 实现的 Adam，称为 DeepSpeedCPUAdam。DeepSpeedCPUAdam 比标准 PyTorch 实现快 5 倍到 7 倍。要深入了解 ZeRO-Offload 的设计和性能，请参阅我们的 博客文章。

训练环境

在本教程中，我们将使用 DeepSpeed Megatron-LM GPT-2 代码配置一个 100 亿参数的 GPT-2 模型。如果您之前尚未执行此操作，我们建议您逐步完成 Megatron-LM 教程。我们将使用一个具有 32GB RAM 的 NVIDIA Tesla V100-SXM3 Tensor Core GPU 用于此练习。

在单个 V100 GPU 上训练一个 100 亿参数的 GPT-2 模型

我们需要更改 Megatron-LM 启动脚本和 DeepSpeed 配置 json。

Megatron-LM GPT-2 启动脚本更改

我们需要对 DeepSpeed Megatron-LM GPT-2 模型的启动脚本进行两处更改。第一个更改是配置一个启用激活检查点的 100 亿参数 GPT-2 模型，这可以通过以下更改集来实现

       --model-parallel-size 1 \
       --num-layers 50 \
       --hidden-size 4096 \
       --num-attention-heads 32 \
       --batch-size 10 \
       --deepspeed_config ds_zero_offload.config \
       --checkpoint-activations

如果您已完成 Megatron-LM 教程，则上述更改中的大多数标志应该很熟悉。

其次，我们需要应用以下更改以确保仅使用一个 GPU 进行训练。

   deepspeed --num_nodes 1 --num_gpus 1 ...

DeepSpeed 配置更改

ZeRO-Offload 利用了许多 ZeRO 阶段 1 和 2 机制，因此启用 ZeRO-Offload 的配置更改是启用 ZeRO 阶段 1 或 2 所需更改的扩展。启用 ZeRO-Offload 的 zero_optimization 配置如下所示

{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
        }
        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

如上所示，除了将 stage 字段设置为 2（以启用 ZeRO 阶段 2，但阶段 1 也适用）之外，我们还需要将 offload_optimizer 设备设置为 cpu 以启用 ZeRO-Offload 优化。此外，我们可以设置其他 ZeRO 阶段 2 优化标志，例如 overlap_comm 以调整 ZeRO-Offload 性能。通过这些更改，我们现在可以运行模型了。我们在下面分享了一些训练的屏幕截图。

这是训练日志的屏幕截图

这是 nvidia-smi 的屏幕截图，显示在训练期间只有 GPU 0 处于活动状态

最后，这是 htop 的屏幕截图，显示优化器计算期间的主机 CPU 和内存活动

恭喜！您已完成 ZeRO-Offload 教程。

更新：