IBM 公布 PyTorch 的突破性进展，加快 AI 模型训练速度

IBM 研究部宣布 PyTorch 框架取得重大进展，旨在提高 AI 模型训练的效率。这些改进在 PyTorch 大会上进行了展示，重点介绍了能够处理海量数据的新数据加载器以及大型语言模型 (LLM) 训练吞吐量的显著增强。

PyTorch 数据加载器的增强功能

新的高吞吐量数据加载器允许 PyTorch 用户在多台机器上无缝分配 LLM 训练工作负载。这项创新使开发人员能够更有效地保存检查点，从而减少重复工作。据 IBM Research 称，该工具是 Davis Wertheimer 和他的同事出于需要而开发的，他们需要一种解决方案来有效地管理和在多台设备上传输大量数据。

最初，团队面临着现有数据加载器的挑战，这导致训练过程出现瓶颈。通过迭代和改进方法，他们创建了一个支持动态和适应性操作的 PyTorch 原生数据加载器。此工具可确保不会重新访问之前看到的数据，即使资源分配在作业中途发生变化。

在压力测试中，数据加载器在一个月的连续运行中成功传输了 2 万亿个令牌，没有出现任何故障。它展示了每个工作者每秒加载超过 90,000 个令牌的能力，相当于 64 个 GPU 上每天加载 5 万亿个令牌。

最大化训练吞吐量

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

这些进步共同旨在将更快的基于云的模型训练从实验阶段转移到更广泛的社区应用，从而有可能改变人工智能模型训练的格局。

图片来源：Shutterstock