ZeRO++

ZeRO++ 是一个基于 ZeRO 之上的通信优化策略系统,无论模型规模或跨设备带宽限制如何,它都能为大型模型训练提供无与伦比的效率。阅读我们的 ZeRO++ 博客论文 了解更多信息!

我们建议您在学习本教程之前阅读关于 入门ZeROMegatron-DeepSpeed 的教程。

ZeRO++ 的三个组成部分

ZeRO++ 包含三个关键设计,即量化权重 (qwZ)、分层分区 ZeRO (hpZ) 和量化梯度 (qgZ)

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

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

训练环境

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

使用 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++ 教程。

更新: