当你在海量数据上训练人工智能模型时,如何克服瓶颈?在今年的 PyTorch 会议上,IBM Research 展示了一款用于大规模 LLM 培训的突破性数据加载器。该工具现已可供 PyTorch 用户使用,旨在为尽可能广泛的受众简化大规模培训。
研究的起源
高吞吐量数据加载器的想法源于研究科学家在模型训练期间观察到的实际问题,因为他们的工作需要一种可以跨多个设备处理大量数据的工具,同时跟上日益高效的 GPU 的步伐。正如 IBM Research 在其关于该版本的博客中指出的那样,“这一切都归功于一组研究人员,他们只是构建了完成工作所需的工具。”
IBM 研究中心的 Davis Wertheimer 解释了大规模培训期间可能出现的一些挑战:“大规模培训存在 80/20 规则。百分之八十的已发表文献都在研究 GPU 内存与通信和计算之间的算法权衡。但当你真正尝试构建某些东西时,80% 的时间,你可以依赖所有这些其他实际问题的长尾,因为管道以最窄瓶颈的速度运行。”
当 IBM 团队开发培训平台时,他们不断遇到瓶颈。 “随着我们越来越擅长使用 GPU,瓶颈越来越多地出现在数据加载器上,”Wertheimer 说道。
这种认识导致了双重开发过程。 “这是一个并行的过程,一方面,不断发展我们的培训平台,另一方面,不断发展我们的数据加载器,以满足我们培训平台的速度需求,避免出现瓶颈,”他解释道。
世界一流的数据加载器的主要特点
IBM 研究中心的 Linsong Chu 概述了数据加载器的基本特性:
有状态和可检查点: “每当你保存模型时,你的数据加载器状态也会被保存,每当你从检查点恢复时,模型状态和数据加载器状态都需要同时恢复,”Chu 说。
自动调整检查点: 数据加载器会在扩展训练期间自动调整以适应工作负载的变化。 Chu 指出:“培训很容易需要数周或数月的时间,并且有很多原因导致您可能必须在中间重新调整工作量。”
高效的数据流: 该系统支持数据流,并且洗牌的构建开销为零。
异步分布式操作: “我们希望数据加载器是非阻塞的,”Chu 解释道。 “在保存数据加载器状态时,我们希望以零通信的形式分发保存内容。”
动态数据混合: 数据加载器可以适应不同的数据混合比例,这对于不断变化的训练需求非常有用。
高效的全局洗牌: 该工具解决了处理大型数据集时的内存瓶颈,即使数据增长时也能提高洗牌效率。
PyTorch 原生、模块化且广泛: 数据加载器专为适应性和可扩展性而设计,为未来的增长做好了准备。 “如果明年我们必须处理 30 万亿、50 万亿或 100 万亿代币怎么办?”楚问道。 “世界瞬息万变,因此我们需要构建数据加载器,使其不仅能够生存在今天,而且还能生存在明天。”
真实世界的表现
IBM 研究团队在几个月内严格测试了他们的数据加载器,运行了数百个大大小小的作业。他们观察到稳定且流畅的代码。而且,整个数据加载器异步运行并且是非阻塞的。
“我们利用了许多内置的 PyTorch 功能来实现这一切,”Wertheimer 说。 “这就是我们做出贡献、回馈社会的原因。”
电子书:如何选择合适的基础模型
这篇文章有帮助吗?
是的不