ZeRO++

ZeRO++ 是一套建立在 ZeRO 之上的通信优化策略系统，旨在为大规模模型训练提供无与伦比的效率，无论规模或跨设备带宽限制如何。阅读我们的 ZeRO++ 博客 和 论文 以了解更多信息！

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

ZeRO++ 的三个组件

ZeRO++ 包含三个关键设计，即量化权重（qwZ）、分层分区 ZeRO（hpZ）和量化梯度（qgZ）：

  • qwZ 应用基于块的量化，将ZeRO参数的全收集通信量从FP16减少到INT8，减少一半。
  • hpZ 通过数据重映射和重新计算消除了节点间的反向参数全收集通信。
  • qgZ 使用一种新的通信效率高的基于全对全的量化梯度平均来替换梯度全归约集合。

总的来说，与ZeRO基线相比，这三种优化将通信量减少了4倍。这三种组件中的每一种都可以独立启用，也可以作为一个组集体启用，如下一节所述。

训练环境

在本教程中，我们将使用DeepSpeed的Megatron-DeepSpeed GPT-2代码配置一个180亿参数的GPT-2模型。我们将使用4个节点，每个节点配备16个NVIDIA Tesla V100-SXM3 Tensor Core GPU，每个节点具有32GB内存来进行此练习。

使用ZeRO++训练一个18B参数的GPT-2

用户代码无需更改。然而，由于ZeRO++扩展了ZeRO Stage 3（ZeRO-3），需要添加适当的标志来激活每个或所有三个ZeRO++通信集体优化。以下是三个标志及其含义、默认值和推荐值：

  • zero_quantized_weights: 布尔值，指示是否使用量化零权重（qwZ），默认值为 false。
  • zero_hpz_partition_size: hpZ（次要分区）组中的排名数量，默认值为1表示没有hpZ，理想值是每个节点的排名（GPU）数量。
  • zero_quantized_gradients: 布尔值，指示是否使用量化零梯度（qgZ），默认值为 false。

DeepSpeed 配置更改

下面展示了一个启用了所有三种ZeRO++优化的deepspeed配置示例片段：

{
    "zero_optimization": {
        "stage": 3,
        "reduce_bucket_size": 10000000,
        "reduce_scatter": true,

        "zero_quantized_weights": true,
        "zero_hpz_partition_size": 16,
        "zero_quantized_gradients": true,

        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

最后，要启动您的实验，请发出以下命令：

       deepspeed pretrain_zeropp_gpt.py \
       --tensor-model-parallel-size 1 \
       --pipeline-model-parallel-size 1 \
       --num-layers 40 \
       --hidden-size 6144 \
       --seq-length 512 \
       --num-attention-heads 32 \
       --batch-size 1 \
       --zero-stage 3 \
       --deepspeed_config ds_zeropp_config.json \
       --deepspeed-activation-checkpointing \
       --fp16 \
       --checkpoint-activations

查看更多关于Megatron-DeepSpeed的详细信息，请参阅教程示例，了解如何启动Megatron-DeepSpeed任务。

以下是ZeRO基线和ZeRO++的训练日志截图：

ZeRO 基线

ZeRO++

恭喜！您已经完成了ZeRO++教程。

更新: