ZeRO-卸载

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

我们建议您在开始本教程之前，先阅读入门和ZeRO的教程。

ZeRO-Offload 是一种 ZeRO 优化技术，它将优化器的内存和计算从 GPU 转移到主机 CPU。ZeRO-Offload 使得具有多达 130 亿参数的大型模型能够在单个 GPU 上高效训练。在本教程中，我们将使用 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内存的NVIDIA Tesla V100-SXM3 Tensor Core GPU来进行此练习。

在单个V100 GPU上训练一个10B参数的GPT-2

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

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

我们需要对DeepSpeed Megatron-LM GPT-2模型的启动脚本进行两项更改。第一项更改是配置一个启用了激活检查点的10B参数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教程。

更新: