解决多源大模型训练中数据加载的负载不均与内存浪费问题
MegaScale-Data: Scaling Dataloader for Multisource Large Foundation Model Training
- 分离数据预处理角色,减少重复读取和内存占用
- 实现4.5倍训练吞吐提升,CPU内存使用降低13.5倍
- 适合需要动态混合多源数据的大模型训练场景
当前大模型训练框架采用数据并行方式加载数据,每个数据加载器处理独立的数据子集。当训练数据来自多个异构源时,面临两大挑战:一是注意力机制的二次复杂度导致数据并行端点间负载不均,降低训练效率;二是各数据源需独立维护文件访问状态,造成冗余内存开销,并阻碍动态数据混合(如课程学习)及混合并行下的高效访问。本文提出MegaScale-Data,一种工业级多源大模型训练数据加载架构,包含三项核心创新:(1) 通过角色分离的数据预处理(源加载器/数据构造器),消除源与并行间的冗余数据访问,保障多源可扩展性;(2) 中心化声明式数据平面,支持训练时的长短期上下文、多模态及课程学习等动态调度;(3) 针对异构预处理成本的多层次自动分区与扩展机制。同时分享部署与容错设计经验。实验表明,该方案实现最高4.5倍端到端训练吞吐提升,以及13.5倍的CPU内存消耗降低。
原文摘要 · Abstract (English)
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. When preparing data for LFM training that originates from multiple, distinct sources, two fundamental challenges arise. 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 MegaScale-Data, an industrial-grade distributed data loading architecture for multisource LFMs training, with three key 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 load-time 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. We also contribute our designs and operational experience in deployment and fault tolerance. MegaScale-Data achieves up to: (1) 4.5x end-to-end training throughput improvement, and (2) 13.5x reduction in CPU memory usage.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。