MegaScale-Data: Scaling Dataloader for Multisource Large Foundation Model Training

本文提出了 MegaScale-Data，一种专为多源大基础模型训练设计的工业级分布式数据加载架构，通过解耦数据预处理、集中式数据编排及多级自动分区机制，有效解决了多源数据并行训练中的负载不均衡与内存冗余问题，实现了高达 4.5 倍的端到端训练吞吐量提升和 13.5 倍的 CPU 内存占用降低。

要点

这篇论文介绍了一个名为 MegaScale-Data 的新技术，它的核心任务是解决在训练超大型人工智能模型（比如能写诗、画图的“大模型”）时，如何高效、省心地准备和输送数据的问题。

为了让你更容易理解，我们可以把训练大模型想象成经营一家超级繁忙的“米其林餐厅”。

1. 背景：餐厅的困境（为什么要做这个？）

想象一下，你的餐厅（AI 模型）要接待成千上万的客人（训练数据）。这些客人来自四面八方：有的来自“文字村”（文本数据），有的来自“图片镇”（图像数据），还有的来自“视频城”（视频数据）。

在传统的做法中，餐厅里有很多个备餐员（DataLoader，数据加载器），每个备餐员负责一部分客人。但是，当客人来源太杂、数量太多时，会出现两个大麻烦：

• 麻烦一：忙闲不均（负载不平衡）
  • 比喻：有的客人只点了一杯水（短文本），备餐员几秒钟就搞定；有的客人点了一桌满汉全席（长文本或高清图片），备餐员要忙活半小时。
  • 后果：因为餐厅是流水线作业，所有备餐员必须等最慢的那一个做完，大家才能一起上菜。结果就是，做快餐的备餐员在干等，做满汉全席的备餐员累得半死，整个餐厅效率极低。
• 麻烦二：冰箱塞爆了（内存浪费）
  • 比喻：每个备餐员为了应对不同的客人，都在自己的小冰箱里存了一套完整的“菜单、食材清单、甚至切菜工具”。如果有 100 种不同的食材来源，每个备餐员都要存 100 套工具。
  • 后果：餐厅的冰箱（CPU 内存）很快就被这些重复的工具塞满了，根本放不下真正的食材。而且，如果客人突然想换菜单（动态调整数据比例），每个备餐员都要重新整理冰箱，乱作一团。

2. 解决方案：MegaScale-Data 的“新厨房”架构

MegaScale-Data 就像是一个智能厨房管理系统，它彻底改变了备餐的方式，提出了三个核心创新：

创新一：拆分工种（Disaggregated Preprocessing）

• 旧模式：每个备餐员既是“切菜工”又是“摆盘工”，还要负责“点菜”。
• 新模式：
  • 专职采购员（Source Loaders）：他们只负责去特定的“菜市场”（数据源）把菜买回来，洗好，切成块。不管客人是谁，他们只负责把“原材料”准备好。
  • 专职摆盘师（Data Constructors）：他们坐在流水线末端，专门负责把采购员送来的菜，根据客人的具体要求（比如“要长盘子”或“要短盘子”）进行最后的组装和打包。
• 好处：采购员不用管怎么摆盘，摆盘师不用去菜市场。大家各司其职，不再重复购买工具，冰箱空间瞬间释放了。

创新二：中央调度室（Centralized Data Plane）

• 旧模式：每个备餐员自己决定今天先做谁的菜，容易乱套。
• 新模式：有一个超级调度员（Planner）。他手里有一张巨大的“智能地图”（DGraph），上面清楚标记着：
  • 哪些菜是“短菜”（短文本），哪些是“长菜”（长文本）。
  • 哪些客人（GPU 显卡）正在等菜，哪些已经吃饱了。
  • 动态调整：如果今天“文字村”的客人变多了，调度员会立刻指挥采购员多去文字村进货；如果“图片镇”的客人变少了，就减少那边的采购。
• 好处：就像交通指挥中心一样，确保所有流水线上的工人（GPU）都能均匀地拿到工作量，没人闲着，也没人累死。

创新三：弹性伸缩（Auto-Scaling）

• 旧模式：餐厅不管客人多还是少，都固定雇佣 100 个备餐员。客人少时浪费钱，客人多时忙不过来。
• 新模式：系统能自动数人头。
  • 如果“视频城”的客人突然变多（处理视频很费时间），系统会自动增加专门处理视频的采购员。
  • 如果“文字村”的客人变少，就减少那边的采购员。
• 好处：根据实际工作量动态调整人手，既省钱又高效。

3. 效果如何？（成绩单）

这套新系统在实际的大规模测试中（比如在 4096 块显卡的超级集群上）表现惊人：

1. 速度快了 4.5 倍：餐厅上菜的速度（训练吞吐量）提升了 4.5 倍，因为不再有人干等，也没有人因为冰箱塞满而手忙脚乱。
2. 省了 13.5 倍的内存：因为不再让每个备餐员都重复存储工具，冰箱（内存）的使用量大幅降低，这意味着可以用更少的钱办更大的事。
3. 更聪明：即使客人突然改变口味（动态调整数据混合比例），或者某个备餐员突然生病（故障），系统也能立刻找到替补，保证餐厅不停业。

总结

MegaScale-Data 就是把原本混乱、重复、低效的“大模型数据准备”过程，变成了一条分工明确、中央指挥、灵活应变的现代化智能流水线。

它不再让每个 AI 训练节点都“单打独斗”地处理数据，而是通过拆解任务和统一调度，让成千上万个计算芯片能像一支训练有素的交响乐团一样，整齐划一地演奏出最强大的 AI 模型。

技术摘要

MegaScale-Data 技术总结

1. 研究背景与问题 (Problem)

随着大型基础模型（LFM，如大语言模型和视觉 - 语言模型）的兴起，训练数据呈现出多源（Multisource）、**异构（Heterogeneous）和动态混合（Dynamic Mixing）**的特征。现有的数据加载框架（Dataloaders）在处理此类场景时面临两大核心挑战：

1. 工作负载不平衡（Workload Imbalance）：
   • 原因： 注意力机制（Attention Operator）具有 $O(l^2)$ 的二次计算复杂度。由于不同数据源（如文本、图像、视频）的序列长度分布不均，导致数据并行（Data Parallel）各 Rank 之间的计算负载严重失衡。
   • 后果： 产生“拖尾”（Stragglers）现象，导致 GPU 等待，降低训练效率。此外，在多模态模型（如 VLM）中，不同模块（如视觉编码器与语言骨干网）对同一批次数据的处理需求不同，进一步加剧了模块间的不平衡。
2. 内存开销过大与扩展性受限（Memory Overhead & Scalability）：
   • 原因： 传统架构中，每个并行加载器（Loader）都独立维护所有数据源的文件访问状态（如 Socket 连接、元数据、I/O 缓冲区）。当数据源数量增加（数百至数千个）或并行度（如 Pipeline Parallelism, Context Parallelism）扩大时，内存开销呈线性甚至更剧烈增长。
   • 后果： 造成巨大的内存浪费，限制了混合并行策略下的扩展能力。同时，缺乏协调机制导致在混合并行（如 PP/CP）中，多个 GPU 重复获取和预处理相同的数据，造成 I/O 和内存的冗余。
3. 动态调度困难：
   • 课程学习（Curriculum Learning）或基于实时指标（如 Loss）的动态数据混合策略要求数据加载器能够灵活调整采样比例，而现有系统缺乏对此类动态策略的原生支持。

2. 方法论 (Methodology)

MegaScale-Data 提出了一种工业级分布式数据加载架构，旨在解决多源 LFM 训练中的可扩展性和编排问题。其核心设计理念是**“解耦（Disaggregation）”与“集中式编排（Centralized Orchestration）”**。

2.1 核心架构组件

系统基于 Actor 模型构建，主要包含三个核心组件：

• Source Loaders（源加载器）： 专门负责特定数据源的样本级转换（如 JPEG 解码、文本分词）。它们作为独立的 Actor 运行，消除了跨并行 Rank 的文件访问状态冗余。
• Data Constructors（数据构造器）： 作为数据并行组的汇聚点，负责聚合 Source Loaders 的输出，执行批次级操作（如 Tensor 填充、打包）和并行变换（Parallelism Transformation）。它们支持跨 Rank 的数据共享，避免重复加载。
• Planner（规划器）： 系统的“大脑”，负责全局数据编排。它包含两个关键子模块：
  • DGraph： 有向无环图，用于追踪样本生命周期、依赖关系和数据流状态。
  • ClientPlaceTree： 逻辑设备拓扑模型，感知混合并行策略（DP/PP/CP/TP），指导数据如何分发给不同的训练节点。

2.2 关键机制

1. 声明式数据编排（Declarative Data Orchestration）：
   • 提供了一套高级原语（Primitives），如 mix（动态混合）、distribute（并行分发）、cost（成本建模）、balance（负载均衡）和 broadcast_at（选择性广播）。
   • 用户无需编写底层执行代码，即可定义复杂的多源混合策略（如长 - 短序列混合、多模态对齐）和负载均衡逻辑。
2. 多级自动分区与弹性扩缩容（Multi-level Auto-partitioning & Scaling）：
   • 离线自动分区： 根据数据源的异构处理成本（如图像解码比文本分词慢得多），将数据源自动划分为不同数量的 Source Loader Actor，优化资源分配。
   • 在线混合驱动扩缩容： 根据训练过程中的动态采样权重变化，Planner 实时调整 Source Loaders 的数量和分片策略，确保数据供给速率与训练需求匹配，避免 GPU 空闲。
3. 混合并行感知的冗余消除：
   • 针对 Pipeline Parallelism (PP) 和 Context Parallelism (CP)，系统通过共享预处理的样本和元数据，消除了不同 Rank 间重复的数据获取和预处理开销。
4. 高可用容错机制：
   • 引入**Shadow Loaders（影子加载器）**作为热备，结合差异化的检查点（Differential Checkpointing）机制，在节点故障时实现快速故障转移，保证数据流不中断。

3. 主要贡献 (Key Contributions)

1. 解耦的多源预处理架构： 设计了基于 Actor 模型的分布式预处理流水线，通过 Source Loaders 和 Data Constructors 的角色分离，彻底消除了多源数据加载中的源级冗余和并行级冗余。
2. 集中式声明式数据平面： 提出了 DGraph 和 ClientPlaceTree 抽象，实现了混合并行感知的数据编排，支持长 - 短序列、多模态和课程学习等复杂策略的声明式定义，极大降低了开发复杂度。
3. 自适应多源扩展机制： 提出了基于异构预处理成本和动态混合比例的自动分区与扩缩容算法，动态优化 CPU 资源利用率。
4. 工业级部署经验： 提供了在大规模集群（最高 4096 GPU）上的部署细节、容错策略和弹性重分片（Resharding）机制。

4. 实验结果 (Results)

在 ByteDance 的工业级集群（最高 4096 张 GPU）上，针对视觉 - 语言模型（VLM）和纯文本大模型的训练进行了评估：

• 训练吞吐量提升： 相比数据并行基线，MegaScale-Data 实现了高达 4.5 倍 的端到端训练吞吐量提升。在长上下文（16k tokens）场景下，提升尤为显著（最高达 4.54 倍）。
• 内存使用降低： 通过消除冗余的文件访问状态和并行加载器实例，CPU 内存使用量减少了 13.5 倍。
• 扩展性： 在扩展到 4096 GPU 时，系统保持了良好的扩展性，而传统基线在大规模下因通信瓶颈和内存溢出而崩溃。
• 开销分析： 数据获取延迟完全被训练计算掩盖，Planner 的编排开销极小（毫秒级），且随集群规模增长可控。
• 训练收敛性： 负载均衡策略对训练 Loss 的收敛几乎没有负面影响，仅在引入 Context Parallelism 时产生微小波动，但不影响最终收敛。

5. 意义与价值 (Significance)

• 解决 LFM 训练瓶颈： 针对当前大模型训练中日益突出的“数据效率”问题，MegaScale-Data 证明了优化数据加载和编排是提升训练效率的关键，其效果不亚于模型并行策略的优化。
• 工业级落地范式： 该工作展示了如何在超大规模集群中处理数百个异构数据源的动态混合问题，为未来万亿参数模型的多模态训练提供了可复用的架构范式。
• 资源成本节约： 大幅降低的 CPU 内存占用意味着在同等硬件配置下可以支持更大的 Batch Size 或更多的并发训练任务，直接降低了训练成本。
• 灵活性与通用性： 其声明式接口和自动扩缩容机制，使得训练框架能够轻松适应课程学习、动态数据混合等高级训练策略，无需修改底层代码。

综上所述，MegaScale-Data 通过架构创新和智能编排，成功解决了多源大规模基础模型训练中的数据加载瓶颈，显著提升了训练效率和资源利用率。