ZeRO++

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

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

ZeRO++ 的三个组成部分

ZeRO++ 由三个关键设计组成，即量化权重 (qwZ)、分层分区 ZeRO (hpZ) 和量化梯度 (qgZ)

  • qwZ 应用基于块的量化，将 ZeRO 参数的全收集通信量从 FP16 减少到 INT8，减少一半。
  • hpZ 通过数据重映射和重新计算，消除了节点间反向传播参数的全收集通信。
  • qgZ 用新的基于 all-to-all 的通信高效量化梯度平均来取代梯度 allreduce 集合操作。

总的来说，这三项优化将通信量与 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（二级分区）组中的排名（rank）数量，默认为 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++ 教程。

更新时间：