混合精度 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++ 相比，hpZ 在混合精度 ZeRO++ 中适用于反向传播和生成过程。

总的来说，这些优化为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++教程。

更新: