混合精度 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++ 教程。

更新时间：