IBM发布PyTorch的重要突破提升AI模型训练速度

IBM研究院宣布了在PyTorch框架中的重大进展，旨在提高AI模型训练的效率。这些改进在PyTorch会议上发布，重点介绍了一种能够处理海量数据的新型数据加载器和对大型语言模型（LLM）训练吞吐量的显著增强。

增强PyTorch的数据加载器

新的高吞吐量数据加载器允许PyTorch用户无缝地将LLM训练工作负载分布到多台机器上。这个创新使开发人员能够更高效地保存检查点，从而减少重复工作。根据IBM研究院的说法，这个工具是由Davis Wertheimer及其同事开发出来的，他们需要一种解决方案来高效地管理和流式处理跨多个设备的海量数据。

起初，团队在使用现有数据加载器时面临挑战，这些加载器在训练过程中造成瓶颈。通过反复迭代和优化，他们创建了一个原生的PyTorch数据加载器，支持动态和可适应的操作。该工具确保即使在资源分配中途改变的情况下，也不会重新访问之前看到的数据。

在压力测试中，数据加载器在一个月的连续运行中成功流式传输了2万亿个标记，没有任何故障。它展示了每个工作者每秒加载超过9万个标记的能力，相当于在64个GPU上一天加载半万亿个标记。

最大化训练吞吐量

IBM研究的另一个重要关注点是优化GPU使用，以防止AI模型训练中的瓶颈。团队使用完全分片的数据并行（FSDP）技术将大型训练数据集均匀分布在多台机器上，从而提高了模型训练和调优的效率和速度。将FSDP与torch.compile结合使用，带来了显著的吞吐量提升。

IBM研究科学家Linsong Chu指出，他们的团队是最早使用torch.compile和FSDP训练模型的团队之一，使用A100 GPU实现了每秒每GPU 4,550个标记的训练速度。这一突破通过最近在Red Hat Enterprise Linux AI（RHEL AI）上发布的Granite 7B模型得到了验证。

还在探索进一步的优化，包括集成Nvidia H100 GPU支持的FP8（8位浮点）数据类型，这已显示出高达50%的吞吐量增益。IBM研究科学家Raghu Ganti强调了这些改进在基础设施成本降低方面的显著影响。

未来前景

IBM研究继续探索新领域，包括在IBM的人工智能单元（AIU）上使用FP8进行模型训练和调优。团队还专注于Nvidia的开源AI部署和执行软件Triton，该软件旨在通过将Python代码编译到特定的硬件编程语言中进一步优化训练。

这些进展共同旨在将云端基于实验阶段的模型训练加速到更广泛的社区应用，可能会改变AI模型训练的格局。