混合精度 ZeRO++

混合精度 ZeRO++ (MixZ++) 是一套基于 ZeRO 和 ZeRO++ 的优化策略，旨在提高效率并减少大型模型训练和推理的内存使用，当用户使用 低秩自适应 (LoRA) 训练时。MixZ++ 将模型参数跨 GPU 分区以减少占用空间，并且仅在需要时使用量化通信收集它们，类似于其 ZeRO 和 ZeRO++ 兄弟。我们的评估表明，对于在 128 个 V100 GPU 上运行的 Llama-2-70B 模型，MixZ++ 将训练吞吐量提高了高达 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++ 教程。

更新：