Orchestrate Multimodal Data with Batch Post-Balancing to Accelerate Multimodal Large Language Model Training

📄 arXiv: 2503.23830v2 📥 PDF

作者: Yijie Zheng, Bangjun Xiao, Lei Shi, Xiaoyang Li, Faming Wu, Tianyu Li, Xuefeng Xiao, Yang Zhang, Yuxuan Wang, Shouda Liu

分类: cs.DC, cs.AI

发布日期: 2025-03-31 (更新: 2025-04-09)


💡 一句话要点

OrchMLLM框架通过批量后平衡加速多模态大语言模型训练

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 多模态大语言模型 训练加速 数据并行 批量后平衡 模态组成不一致性

📋 核心要点

  1. 现有MLLM训练受限于模态组成不一致性,导致小批量数据不平衡,GPU利用率低,训练效率低下。
  2. OrchMLLM框架通过批量后平衡调度器消除小批量不平衡,并使用全局协调器管理多模态数据。
  3. 实验表明,OrchMLLM在训练84B MLLM时,MFU达到41.6%,吞吐量比Megatron-LM提升高达3.1倍。

📝 摘要(中文)

多模态大语言模型(MLLM),如GPT-4o,正受到广泛关注。在MLLM训练的探索过程中,我们发现模态组成不一致性问题,即不同样本中特定模态的比例差异巨大。这加剧了解决小批量不平衡的挑战,导致数据并行(DP)实例之间的GPU利用率不均,严重降低了MLLM训练的效率和可扩展性,最终影响训练速度并阻碍对MLLM的进一步研究。为了应对这些挑战,我们引入了OrchMLLM,一个旨在缓解由模态组成不一致性引起的MLLM训练低效率的综合框架。首先,我们提出批量后平衡调度器,一种有效消除顺序数据中小批量不平衡的技术。此外,我们将MLLM全局协调器集成到训练框架中,以协调多模态数据并解决由模态组成不一致性引起的问题。我们在各种MLLM规模上评估了OrchMLLM,证明了其效率和可扩展性。实验结果表明,在使用2560个H100 GPU训练具有三种模态的84B MLLM时,OrchMLLM实现了41.6%的模型FLOPs利用率(MFU),在吞吐量方面优于Megatron-LM高达3.1倍。

🔬 方法详解

问题定义:论文旨在解决多模态大语言模型(MLLM)训练过程中由于“模态组成不一致性”导致的训练效率低下问题。具体来说,不同训练样本中不同模态(如图像、文本、音频)的比例差异很大,这导致在数据并行训练时,各个GPU上的小批量数据负载不均衡,部分GPU空闲,造成资源浪费。现有方法难以有效解决这种小批量不平衡问题,限制了MLLM的训练速度和可扩展性。

核心思路:论文的核心思路是通过“批量后平衡”来解决小批量不平衡问题,并引入全局协调器来管理多模态数据。批量后平衡是指在数据加载后,对一个小批量内的数据进行重新排列,使得每个GPU上的数据在模态组成上尽可能平衡。全局协调器则负责整体的数据调度和模态管理,确保训练过程的稳定和高效。这样设计的目的是为了充分利用每个GPU的计算资源,提高训练效率。

技术框架:OrchMLLM框架主要包含两个核心模块:批量后平衡调度器(Batch Post-Balancing Dispatcher)和MLLM全局协调器(MLLM Global Orchestrator)。批量后平衡调度器负责在每个训练迭代中,对小批量数据进行重新排列,以消除模态组成不平衡。MLLM全局协调器则负责整体的数据调度、模态管理和训练流程控制。整个训练流程如下:首先,数据加载器加载原始数据;然后,批量后平衡调度器对数据进行重新排列;接着,数据被分发到各个GPU进行训练;最后,MLLM全局协调器收集训练结果并进行模型更新。

关键创新:论文最重要的技术创新点是批量后平衡调度器。与传统的静态数据划分方法不同,批量后平衡调度器能够动态地调整小批量数据的模态组成,从而有效地消除小批量不平衡。这种动态调整机制使得每个GPU都能获得相对均衡的数据负载,从而提高GPU利用率。此外,MLLM全局协调器的引入也简化了多模态数据的管理和调度,提高了训练的整体效率。

关键设计:批量后平衡调度器的关键设计在于如何高效地对小批量数据进行重新排列。论文采用了一种基于排序的策略,首先根据每个样本的模态组成计算一个权重,然后根据权重对样本进行排序,最后将排序后的样本分配到各个GPU。MLLM全局协调器的关键设计在于如何有效地管理多模态数据。论文采用了一种基于元数据的管理方法,为每个样本维护一个元数据,记录其模态组成、数据来源等信息。通过元数据,MLLM全局协调器可以方便地进行数据调度和模态管理。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,OrchMLLM在训练具有三种模态的84B MLLM时,在2560个H100 GPU上实现了41.6%的模型FLOPs利用率(MFU),相比于Megatron-LM,吞吐量提升高达3.1倍。这一结果表明OrchMLLM能够显著提高MLLM的训练效率和可扩展性,为更大规模MLLM的训练提供了可能。

🎯 应用场景

该研究成果可广泛应用于各种多模态大语言模型的训练场景,尤其是在数据模态分布不均衡的情况下。例如,在训练包含图像、文本和音频数据的视频理解模型时,OrchMLLM可以显著提高训练效率。此外,该框架还可以应用于其他需要处理多模态数据的任务,如多模态对话系统、多模态检索等。该研究的实际价值在于降低了MLLM的训练成本,加速了MLLM的研究和应用。

📄 摘要(原文)

Multimodal large language models (MLLMs), such as GPT-4o, are garnering significant attention. During the exploration of MLLM training, we identified Modality Composition Incoherence, a phenomenon that the proportion of a certain modality varies dramatically across different examples. It exacerbates the challenges of addressing mini-batch imbalances, which lead to uneven GPU utilization between Data Parallel (DP) instances and severely degrades the efficiency and scalability of MLLM training, ultimately affecting training speed and hindering further research on MLLMs. To address these challenges, we introduce OrchMLLM, a comprehensive framework designed to mitigate the inefficiencies in MLLM training caused by Modality Composition Incoherence. First, we propose Batch Post-Balancing Dispatcher, a technique that efficiently eliminates mini-batch imbalances in sequential data. Additionally, we integrate MLLM Global Orchestrator into the training framework to orchestrate multimodal data and tackle the issues arising from Modality Composition Incoherence. We evaluate OrchMLLM across various MLLM sizes, demonstrating its efficiency and scalability. Experimental results reveal that OrchMLLM achieves a Model FLOPs Utilization (MFU) of $41.6\%$ when training an 84B MLLM with three modalities on $2560$ H100 GPUs, outperforming Megatron-LM by up to $3.1\times$ in throughput.