ZeRO（Zero Redundancy Optimizer）是一种高效的内存管理技术，它通过消除数据并行处理过程中的内存冗余，使得大规模模型训练成为可能。ZeRO技术通过将模型状态（包括优化器状态、梯度以及参数）分割而不是复制到数据并行处理进程中，从而提高了内存效率。 ZeRO有三个阶段：
- ZeRO Stage 1：优化器状态（例如Adam优化器的32位权重、第一和第二时刻估计）被分割到各个进程中，每个进程只更新其分区。
- ZeRO Stage 2：32位梯度的减少（用于更新模型权重）也被分割，每个进程只保留与其优化器状态分区相对应的梯度。
- ZeRO Stage 3：16位模型参数被分割到各个进程中。ZeRO-3会在前向和反向传递期间自动收集和分割它们。 此外，ZeRO-3还包括一个无限卸载引擎，形成ZeRO-Infinity，可以将所有模型状态卸载到CPU和NVMe内存中，从而实现巨大的内存节省。 对于想要深入了解我们算法的用户，请参阅我们关于ZeRO、ZeRO-Offload和ZeRO-Infinity的论文。 请注意，DeepSpeed最初包含的卸载功能是ZeRO-Offload，它可以将优化器和梯度状态卸载到CPU内存中。ZeRO-Infinity是ZeRO-Offload的下一代卸载功能，适用于ZeRO-3。ZeRO-Infinity具有ZeRO-Offload的所有节省，并且能够卸载更多的模型权重，并且具有更有效的带宽利用和计算与通信的重叠。 开始使用： 如果您是DeepSpeed的新手，请查看我们的入门页面。 一旦您开始使用DeepSpeed训练，启用ZeRO-3卸载就像在DeepSpeed配置中启用它一样简单！以下是一些ZeRO-3配置的示例。有关配置和性能调优的完整选项列表，请参阅我们的配置指南。 请注意，ZeRO-Infinity和ZeRO-Offload与我们的深度优化深度speed.ops.adam.DeepSpeedCPUAdam优化器配合使用效果最佳。我们建议使用我们的优化器配置来指示deepspeed.initialize()为您构建优化器。 ZeRO配置： 所有DeepSpeed ZeRO的设置都由DeepSpeedZeroConfig类设置。在主DeepSpeed配置字典的zero_optimization条目下的字典将由这个类解析和验证。参数卸载和优化器卸载的子配置将由DeepSpeedZeroOffloadParamConfig和DeepSpeedZeroOffloadOptimizerConfig解析。 classdeepspeed.runtime.zero.config.DeepSpeedZeroConfig[source] 设置ZeRO优化的参数。 stage: ZeroStageEnum= 0 选择不同阶段的ZeRO优化器。第0、1、2和3阶段分别指的是禁用、优化器状态分区、优化器+梯度状态分区以及优化器+梯度+参数分区。 contiguous_gradients: bool= True 在产生梯度时将它们复制到一个连续的缓冲区。避免在反向传递期间内存碎片化。 reduce_scatter: bool= True 使用reduce或reduce scatter而不是allreduce来平均梯度 reduce_bucket_size: int= 500,000,000 一次减少/全减少的元素数量。限制在大型模型大小下所需的内存 Constraints minimum = 0 use_multi_rank_bucket_allreduce: bool= True 结合不同等级的减少桶并执行All-Reduce而不是多个Reduce操作。当模型较小时，此功能对于在太多GPU上扩展非常有用，从而减少每个数据包的消息大小。 allgather_partitions: bool= True 在每个步骤结束时从所有GPU收集更新后的参数，选择在所有聚集集体之间进行选择，还是在广播集体之间进行选择 allgather_bucket_size: int= 500,000,000 一次全聚集的元素数量。限制在大型模型大小下所需的内存 Constraints minimum = 0 overlap_comm: bool= None 尝试将梯度的减少与反向计算重叠 load_from_fp32_weights: bool= True 一个布尔值，指示是否从检查点中的fp32主权重初始化fp32主权重（无精度损失）或从模型的fp16副本初始化（有精度损失）。这可以用于在检查点缺少优化器状态时初始化优化器状态。 elastic_checkpoint: bool= False 启用加载由不同GPU计数的工作保存的检查点。不再支持。 offload_param: Optional[DeepSpeedZeroOffloadParamConfig]= None 启用将模型参数卸载到CPU或NVMe。这为更大的模型或批量大小释放了GPU内存。仅适用于第3阶段。期望包含DeepSpeedZeroOffloadParamConfig值的字典。 offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig]= None 启用将优化器状态卸载到CPU或NVMe以及将优化器计算卸载到CPU。这为更大的模型或批量大小释放了GPU内存。对于Ze ZeRO Stage 1、2、3的有效。期望包含DeepSpeedZeroOffloadOptimizerConfig值的字典。 sub_group_size: int= 1,000,000,000 用于处理参数的瓦片大小，以适应具有万亿参数的大规模模型。由ZeRO3-Offload和ZeRO-Infinity使用 Constraints minimum = 0 cpu_offload_param: bool= None 已弃用，请使用offload_param cpu_offload_use_pin_memory: bool= None 已弃用，请使用offload_param或offload_optimizer cpu_offload: bool= None 已弃用，请使用offload_optimizer prefetch_bucket_size: int= 50,000,000(alias 'stage3_prefetch_bucket_size') 可以提前获取的参数元素的最大数量。由ZeRO3、ZeRO3-Offload、ZeRO-Infinity以及ZeRO-Inference使用。 Constraints minimum = 0 param_persistence_threshold: int= 100,000(alias 'stage3_param_persistence_threshold') 不分区小于此阈值的参数。较小的值使用较少的内存，但可能会大大增加通信（特别是延迟限制的消息）。 Constraints minimum = 0 model_persistence_threshold: int= sys.maxsize(alias 'stage3_model_persistence_threshold') 可以在GPU中驻留并且不被分区的最大参数元素数量。这对param_persistence_threshold设置产生的未分区参数的数量施加了上限。由ZeRO3-Offload、ZeRO-Infinity以及ZeRO-Inference使用。 Constraints minimum = 0 max_live_parameters: int= 1,000,000,000(alias 'stage3_max_live_parameters') 每个GPU上驻留的参数的最大数量，在释放之前。较小的值使用较少的内存，但执行更多的通信。 Constraints minimum = 0 max_reuse_distance: int= 1,000,000,000(alias 'stage3_max_reuse_distance') 如果将在此阈值内的参数重用，请不要释放参数。较小的值使用较少的内存，但执行更多的通信。 Constraints minimum = 0 gather_16bit_weights_on_model_save: bool= False(alias 'stage3_gather_16bit_weights_on_model_save') 在保存模型时通过save_16bit_model()合并权重。由于权重在GPU之间分区，它们不是state_dict的一部分，因此当此选项启用时，函数会自动收集权重，然后保存fp16模型权重。 stage3_gather_fp16_weights_on_model_save: bool= False 已弃用，请使用gather_16bit_weights_on_model_save ignore_unused_parameters: bool= True 模块中的未使用参数在静态网络中可能是意外的，但在动态网络中可能是正常的。这控制是否在检测到未使用参数时终止训练并显示错误消息。默认情况下，此选项设置为True，这意味着忽略未使用参数，并且训练将继续。现在仅在第2阶段使用。 legacy_stage1: bool ZeRO的配置允许用户根据需要调整各种参数，以优化模型的训练过程。通过合理配置这些参数，用户可以在不同的硬件环境下更高效地训练大规模模型。 在DeepSpeed中启用ZeRO-3卸载的配置示例：
{
"zero_optimization": {
"stage": 3,
"allgather_partitions": True,
"allgather_bucket_size": 2e8,
"reduce_scatter": True,
"reduce_bucket_size": 2e8,
"overlap_comm": True,
"contiguous_gradients": True,
"cpu_offload": True,
"cpu_offload_params": False,
"cpu_offload_use_pin_memory": True,
"sub_group_size": 1e9,
"stage3_prefetch_bucket_size": 5e7,
"stage3_param_persistence_threshold": 1e5,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_fp16_weights_on_model_save": True
}
}
在这个配置中，我们设置了ZeRO-3的各种参数，包括卸载到CPU的选项，以及各种通信和内存管理参数。用户可以根据自己的模型大小和硬件资源调整这些参数。 总结来说，ZeRO技术是DeepSpeed中的一个关键特性，它通过智能地管理内存和通信，使得大规模模型训练变得更加高效和可行。通过调整ZeRO的配置，用户可以在不同的训练环境中优化模型的性能。
