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 个节点的 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++ 教程。