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 RAM。

使用 ZeRO++ 训练 180 亿参数的 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++ 教程。

更新：