混合精度 ZeRO++
混合精度 ZeRO++ (MixZ++) 是一套基于 ZeRO 和 ZeRO++ 的优化策略,旨在用户使用 低秩适应 (LoRA) 训练时,提高大型模型训练和推理的效率并减少内存使用。MixZ++ 将模型参数分布到多个 GPU 上以减少内存占用,并仅在需要时通过量化通信进行聚合,这与其 ZeRO 和 ZeRO++ 兄弟技术类似。我们的评估表明,MixZ++ 将在 128 块 V100 GPU 上运行的 Llama-2-70B 模型的训练吞吐量提高了高达 3.3 倍。阅读我们的 DeepSpeed Chat 博客、ZeRO++ 博客 和论文以了解更多信息!
我们建议您在学习本教程之前,阅读关于 入门、ZeRO 和 Megatron-DeepSpeed 的教程。
关键设计
混合精度 ZeRO++ (MixZ++) 继承了 ZeRO++ 的关键设计,即量化权重 (qwZ) 和分层分区 ZeRO (hpZ),但适用性有所不同。
- qwZ 对冻结权重应用基于块的量化,以减少内存使用和全收集通信量。与 ZeRO++ 相比,混合精度 ZeRO++ 中的 qwZ 保持冻结权重为量化状态,因此在运行时没有量化开销,并减少了内存使用。
- hpZ 通过数据重映射和重新计算消除了节点间的参数全收集通信。与 ZeRO++ 相比,混合精度 ZeRO++ 中的 hpZ 适用于反向传播和生成过程。
总的来说,这些优化为 LoRA 训练带来了更好的可扩展性和效率。每个组件都可以独立启用,也可以作为一个组一起启用。
启用混合精度 ZeRO++ (MixZ++)
一个可用的 MixZ++ 示例已准备好,位于 MixZ++ 示例脚本。如果您希望在您的管道中手动启用 MixZ++,请参考以下说明。
DeepSpeed 配置更改
下面展示了启用所有 MixZ++ 优化的 DeepSpeed 配置示例片段。
{
"zero_optimization": {
"stage": 3,
"..."
"zero_quantized_nontrainable_weights": true,
"zero_hpz_partition_size": 16,
"..."
}
}
请注意,对于多节点训练,"zero_hpz_partition_size"
应设置为每个节点上的 GPU 数量。例如,如果您每个节点有 8 个 GPU,则 "zero_hpz_partition_size"
应设置为 8。对于单节点训练,不应设置 "zero_hpz_partition_size"
。
训练脚本更改
如果在 DeepSpeed 初始化时传入 LoRA 模型,DeepSpeed 引擎将识别 LoRA 冻结参数。然而,常见的实现方式是先初始化一个基础模型,然后再将其转换为 LoRA 模型。在这种情况下,用户需要在 LoRA 模型转换后显式调用 DeepSpeed 引擎。这只需要一行代码。下面展示了训练脚本的示例片段。
model, optimizer, _, lr_scheduler = deepspeed.initialize(
model=model,
optimizer=optimizer,
args=args,
config=ds_config,
lr_scheduler=lr_scheduler,
dist_init_required=True)
# ...
# (the custom code to convert base model to LoRA model)
# ...
# call DeepSpeed engine again to identify LoRA frozen parameters
model.optimizer.quantize_nontrainable_params()
# ...
恭喜!您已完成混合精度 ZeRO++ 教程。