IBM Research 数据加载器为开源社区增强 AI 模型训练

当您在大量数据上训练 AI 模型时，如何克服瓶颈？在今年的 PyTorch 会议上，IBM Research 展示了一款用于大规模 LLM 训练的开创性数据加载器。该工具现已可供 PyTorch 用户使用，旨在为尽可能广泛的受众简化大规模训练。

高吞吐量数据加载器的想法源于研究科学家在模型训练过程中观察到的实际问题，因为他们的工作需要一种能够跨多个设备处理大量数据的工具——同时还要跟上日益高效的 GPU 发展步伐。正如 IBM Research 在其关于该版本发布的博客中所指出的：“这完全要归功于一支研究人员团队，他们只是构建了完成工作所需的工具。”

IBM Research 的 Davis Wertheimer 解释了大规模训练中可能出现的一些挑战：“在大规模训练方面，存在某种 80/20 法则。已发表文献的 80% 都在研究 GPU 内存、通信与计算之间的算法权衡。但当你真正尝试构建某些东西时，80% 的时间，你都会依赖于所有这些其他实际问题的漫长解决过程，因为管道的运行速度取决于最窄的瓶颈。”

在 IBM 团队开发其训练平台的过程中，他们持续遇到瓶颈。“随着我们越来越擅长使用 GPU，瓶颈越来越频繁地出现在数据加载器上，”Wertheimer 指出。

这一认识引发了一个双重开发过程。“这是一段并行的旅程：一方面，发展我们的训练平台；另一方面，持续改进我们的数据加载器，以满足训练平台的速度需求，避免形成瓶颈，”他解释说。

IBM Research 的 Linsong Chu 概述了该数据加载器的基本特性：

有状态且可检查点保存：“每当你保存一个模型时，你的数据加载器状态也会被保存；而当你从一个检查点恢复时，模型状态和数据加载器状态需要同时恢复，”Chu 说。

检查点的自动重缩放： 数据加载器能自动适应长时间训练过程中工作负载的变化。“训练很容易持续数周或数月，中途有大量原因可能导致你必须重新调整工作负载规模，”Chu 指出。

高效数据流：该系统支持数据流传输，且混洗操作零构建开销。

异步分布式操作：“我们希望数据加载器是非阻塞式的，”Chu 解释道。“在保存数据加载器状态时，我们希望保存操作以零通信的形式分布式进行。”

动态数据混合：该数据加载器能适应不同的数据混合比例，这对于不断演变的训练需求非常有用。

高效全局混洗：该工具解决了处理大型数据集时的内存瓶颈问题，即使数据增长也能高效进行混洗。

PyTorch 原生、模块化且功能广泛：该数据加载器专为适应性和可扩展性设计，为未来的增长做好了准备。“如果明年我们必须处理 30 万亿、50 万亿甚至 100 万亿个词元呢？”Chu 反问道。“世界变化很快，因此我们需要构建的数据加载器不仅要能满足今天的需求，也要能适应明天的发展。”

IBM Research 团队在数月内对其数据加载器进行了严格测试，运行了数百个大大小小的任务。他们观察到了稳定且流畅的代码运行数据。此外，整个数据加载器以异步方式运行且是非阻塞的。

“我们利用了许多 PyTorch 的内置功能才实现了这一切，” Wertheimer 说。“这也是我们将其贡献出来、回馈（给社区）的原因。”