OVERLORD: Ultimate Scaling of DataLoader for Multi-Source Large Foundation Model Training
作者: Juntao Zhao, Qi Lu, Wei Jia, Borui Wan, Lei Zuo, Junda Feng, Jianyu Jiang, Yangrui Chen, Shuaishuai Cao, Jialing He, Kaihua Jiang, Yuanzhe Hu, Shibiao Nong, Yanghua Peng, Haibin Lin, Xin Liu, Chuan Wu
分类: cs.DC, cs.AI
发布日期: 2025-04-14 (更新: 2025-05-18)
💡 一句话要点
OVERLORD：面向多源大模型训练的数据加载器极致扩展方案
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 数据加载器 大模型训练 数据并行 多源数据 负载均衡
📋 核心要点
- 现有大模型训练的数据加载器在多源数据场景下，存在负载不均衡和内存冗余的问题，严重影响训练效率。
- Omniload通过解耦数据预处理、集中式数据平面、多级自动分区和影子加载器等创新技术，实现高效的数据加载和容错。
- 在生产集群上的实验表明，Omniload能够显著提升训练吞吐量并降低内存占用，验证了其有效性和可扩展性。
📝 摘要（中文）
现代大模型训练框架采用数据并行方式，每个数据加载器处理不相交的训练数据子集。在多源数据预处理下，存在两个根本挑战：一是注意力机制的二次计算复杂度导致数据并行rank上的样本分布不均，造成数据加载器之间显著的负载不平衡，降低训练效率；二是支持多样的数据源需要每个数据集的文件访问状态在并行加载器之间冗余复制，消耗过多内存，阻碍动态数据混合（如课程学习），并在混合并行中导致冗余的访问/内存开销。我们提出了Omniload，一个工业级的用于大模型的分布式数据加载架构，包含四个创新点：(1) 通过角色特定的actor（源加载器/数据构造器）进行解耦数据预处理，消除源和并行冗余数据访问，确保多源可扩展性。(2) 用于弹性多源编排的集中式和声明式数据平面，例如长短上下文、多模态和课程学习。(3) 用于异构预处理成本下源加载器的多级自动分区和扩展机制。(4) 具有差分检查点的影子加载器，用于在不中断工作流的情况下进行故障恢复。Omniload部署在扩展到数千个GPU的生产集群上，实现了：(1) 4.5倍的端到端训练吞吐量提升，(2) 13.5倍的CPU内存使用量减少。
🔬 方法详解
问题定义：论文旨在解决多源大模型训练中数据加载器的性能瓶颈问题。现有方法在数据并行场景下，由于数据分布不均和多源数据访问的冗余性，导致负载不平衡和内存开销过大，严重影响训练效率。特别是注意力机制的二次复杂度进一步加剧了负载不平衡的问题。
核心思路：Omniload的核心思路是通过解耦数据预处理流程，将数据加载和数据构造分离到不同的角色（Source Loader和Data Constructor）中，从而消除数据源和并行计算中的冗余访问。同时，采用集中式数据平面进行统一管理和调度，实现弹性多源编排。
技术框架：Omniload的整体架构包含以下几个主要模块：1) Source Loaders：负责从不同的数据源加载原始数据，并进行初步的预处理。2) Data Constructors：负责将加载的数据构建成模型训练所需的格式，例如生成batch。3) Centralized Data Plane：负责管理数据源、调度任务、监控状态等。4) Multi-level Auto-partitioning：根据异构预处理成本自动对Source Loaders进行分区和扩展。5) Shadow Loaders：提供容错机制，通过差分检查点实现快速故障恢复。
关键创新：Omniload的关键创新在于其解耦的数据预处理架构和集中式数据平面。通过将数据加载和数据构造分离，避免了数据源的冗余访问和并行计算中的负载不平衡。集中式数据平面则提供了统一的管理和调度能力，支持弹性多源编排和动态数据混合。
关键设计：Omniload的关键设计包括：1) Role-Specific Actors：Source Loaders和Data Constructors的角色分离，实现了数据预处理流程的解耦。2) Centralized Data Plane：采用声明式配置，支持灵活的数据源管理和任务调度。3) Multi-level Auto-partitioning：根据数据源的特性和计算资源的情况，自动调整Source Loaders的分区策略。4) Differential Checkpointing：Shadow Loaders采用差分检查点技术，减少了检查点的大小和恢复时间。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Omniload在生产集群上实现了显著的性能提升。端到端训练吞吐量提高了4.5倍，CPU内存使用量减少了13.5倍。这些数据验证了Omniload在处理多源大规模数据训练任务时的有效性和可扩展性。
🎯 应用场景
Omniload适用于各种需要处理多源大规模数据的深度学习训练场景，尤其是在大模型训练领域，例如自然语言处理、计算机视觉和多模态学习。它可以显著提升训练效率，降低资源消耗，并支持更灵活的数据管理和调度策略，加速大模型的研发和部署。
📄 摘要（原文）
Modern frameworks for training large foundation models (LFMs) employ dataloaders in a data-parallel manner, with each loader processing a disjoint subset of training data. Under multisource preprocessing, two fundamental challenges exist. First, due to the quadratic computational complexity of the attention operator, the non-uniform sample distribution over data-parallel ranks leads to significant workload imbalance among dataloaders, degrading the training efficiency. Second, supporting diverse data sources requires per-dataset file access states that are redundantly replicated across parallel loaders, consuming excessive memory. This also hinders dynamic data mixing (e.g., curriculum learning) and causes redundant access/memory overhead in hybrid parallelism. We present Omniload, an industrial-grade distributed data loading architecture for LFMs, with four innovations: (1) Disaggregated data preprocessing via role-specific actors (Source Loaders/Data Constructors) to eliminate source and parallelism redundant data access and ensure multisource scalability. (2) Centralized and declarative data plane for elastic multisource orchestration, such as long-short context, multimodality, and curriculum learning. (3) Multi-level auto-partitioning and scaling mechanism for source loaders under heterogeneous preprocessing costs. (4) Shadow loaders with differential checkpointing for fault recovery without workflow interruption. Deployed on production clusters scaling to multi-thousand GPUs, Omniload achieves: (1) 4.5x end-to-end training throughput improvement, (2) 13.5x reduction in CPU memory usage.