TorchTitan项目中InterleavedZeroBubble调度器的性能问题分析

背景介绍

在TorchTitan深度学习框架中，研究人员发现InterleavedZeroBubble调度器在512-GPU规模实验中表现出异常的性能特征。与传统的1F1B和Interleaved1F1B调度器相比，InterleavedZeroBubble调度器显示出显著更高的内存占用(94.45%)和极低的吞吐量(TPS仅为13)，同时模型利用率(MFU)也降至3.41%的低水平。

问题现象

通过基准测试数据对比可以清晰地看到性能差异：

• 1F1B调度器：内存使用82.46GiB(86.80%)，TPS为100，MFU达到26.52%
• Interleaved1F1B调度器：内存使用72.69GiB(76.52%)，TPS提升至128，MFU为33.88%
• InterleavedZeroBubble调度器：内存使用高达89.73GiB(94.45%)，TPS骤降至13，MFU仅为3.41%

值得注意的是，在测试Zero Bubble调度器时，由于兼容性问题，未使用torch.compile、异步TP和Float8等优化技术。

深入调查

进一步研究发现，问题的表现与激活检查点(Activation Checkpointing, AC)技术的使用密切相关：

在不使用AC的情况下：

• Interleaved1F1B：TPS 304，MFU 4.16%
• InterleavedZeroBubble：TPS 339，MFU 4.65% - 表现反而更好

但在启用完整AC后：

• Interleaved1F1B：TPS降至243，MFU 3.33%
• InterleavedZeroBubble：TPS暴跌至100，MFU仅1.37%

根本原因

经过深入分析，发现问题根源在于PyTorch框架内部的一个实现细节。当使用InterleavedZeroBubble调度器配合激活检查点时，会触发特定的执行模式，导致计算图被分割成大量微小片段，进而引发以下问题：

1. 调度开销显著增加
2. 内存管理效率下降
3. 计算资源利用率降低

解决方案

开发团队提出了一个有效的修复方案，通过调整计算图的构建方式，避免了上述问题的发生。测试结果显示该方案成功恢复了InterleavedZeroBubble调度器的预期性能：

在不使用AC的情况下：

• Interleaved1F1B：TPS 206，MFU 3.02%
• InterleavedZeroBubble：TPS提升至227，MFU 3.33%

在使用完整AC的情况下：

• Interleaved1F1B：TPS 177，MFU 2.59%
• InterleavedZeroBubble：TPS 192，MFU 2.82% - 性能差距显著缩小

技术启示

这一问题的解决过程为深度学习系统优化提供了宝贵经验：

1. 调度算法性能可能高度依赖于底层框架实现细节
2. 高级优化技术(如AC)与特定调度策略的组合需要谨慎评估
3. 性能问题的根本原因可能隐藏较深，需要系统性的排查方法
4. 针对特定工作负载的微调往往能带来显著的性能提升

该问题的成功解决不仅修复了InterleavedZeroBubble调度器的性能缺陷，也为TorchTitan框架的稳定性与可靠性做出了重要贡献。