DeepSpeed ZeRO-3 内存估算

1. 基本情况

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live
model = AutoModel.from_pretrained("bigscience/T0_3B")
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1)

• 硬件设置：1个节点，每节点1个GPU
• 模型：bigscience/T0_3B

• • 总参数量：2783M（约27.83亿）
  • 最大层参数量：65M

2. 内存估算结果

CPU 内存	GPU 内存	offload_param	offload_optimizer	zero_init
70.00GB	0.25GB	cpu	cpu	1
70.00GB	0.25GB	cpu	cpu	0
62.23GB	5.43GB	none	cpu	1
62.23GB	5.43GB	none	cpu	0
0.37GB	46.91GB	none	none	1
15.56GB	46.91GB	none	none	0

3. 参数解释

• offload_param：是否将模型参数卸载到CPU
• offload_optimizer：是否将优化器状态卸载到CPU
• zero_init：是否在初始化时就应用ZeRO优化

4. 结果分析

4.1 卸载策略影响

• 当参数或优化器状态卸载到CPU时，CPU内存使用显著增加，GPU内存使用减少
• 卸载策略可以有效减少GPU内存需求，但会增加CPU内存负担

4.2 zero_init 的影响

• zero_init 的影响只在 offload_param=none, offload_optimizer=none 时明显
• 原因：其他情况下，卸载策略的内存需求掩盖了初始化方式的影响

zero_init 工作原理

• zero_init=1：模型参数在初始化时就被分片并分布
• zero_init=0：模型参数先在一个进程中完全初始化，然后再分片分发

zero_init 对内存的影响

• 主要影响CPU内存使用
• 当 offload_param=none, offload_optimizer=none 时：

• • zero_init=1：CPU需要0.37GB
  • zero_init=0：CPU需要15.56GB

zero_init 选择考虑因素

1. 内存效率：zero_init=1 通常更节省内存
2. 初始化速度：zero_init=1 通常更快
3. 模型一致性：zero_init=0 保证所有进程看到相同的初始模型
4. 大模型支持：对于超大模型，zero_init=1 可能是唯一选择

5. 实际应用建议

1. 根据可用硬件资源选择合适的卸载策略
2. 对于内存受限的系统，特别是处理大模型时，考虑使用 zero_init=1
3. 如需精确复现或模型对初始化敏感，可能需要 zero_init=0
4. 实际应用中，建议尝试不同配置以找到最佳设置

6. 注意事项

• 估算值可能与实际运行时的内存使用有差异
• 实际内存需求可能高于估算，因为还有其他运行时开销
• 建议在规划硬件需求时留有余地
• 可以考虑使用其他内存优化技术，如梯度检查点、混合精度训练等

7. 总结

DeepSpeed ZeRO-3 提供了灵活的内存优化选项，允许在不同的硬件配置下高效训练大型模型。理解各参数的作用及其相互影响对于优化训练过程至关重要。在实际应用中，需要根据具体的模型大小、硬件资源和训练需求来选择最合适的配置。

参考：DeepSpeed使用指南(简略版)-CSDN博客