技术摘要：面向多模态大语言模型训练的异构并行

1. 问题陈述

随着基础模型训练从纯文本 Transformer 演变为涵盖图像、音频、视频和文档的多模态系统，训练图在模块深度、宽度和上下文大小方面变得日益异构。然而，当前的并行化栈（如 Megatron-LM、DeepSpeed 等）通常在整个端到端图上强制实施单一、共享的并行布局（张量并行、上下文并行、流水线并行、数据并行和专家并行）。

这种“一刀切”的方法在编码器与大语言模型（LLM）需求出现分歧时会产生显著的低效：

• 不必要的通信：大型 LLM 通常需要高张量并行（TP）度。强制小型编码器继承这些 TP 度，会在本地编码器计算周围引入不必要的 All-Reduce 集合通信，其带来的开销超过了内存节省。
• 上下文并行不匹配：上下文并行（CP）对于 LLM 管理长序列的注意力内存至关重要。然而，编码器的序列长度受模态分辨率（例如图像块）限制，而非 LLM 的上下文窗口。继承 LLM 的 CP 分片会固定编码器序列的大小，却无法提供内存缓解；而那些秩（ranks）本可以作为独立的数据并行（DP）副本更高效地处理更多样本。
• 流水线约束：流水线并行（PP）通常将编码器限制在 LLM 流水线的第一阶段。这迫使编码器继承 LLM 的 TP 度，可能导致特定秩过载，从而阻碍 LLM 利用更高效的流水线布局。

核心挑战在于，计算、内存和通信的最佳工作点会随着工作负载参数（编码器规模、上下文长度、Token 密度、世界规模）的变化而转移。静态的共享布局无法适应这些变化。

2. 方法论

作者提出了异构并行，这是一种运行时抽象，允许单个端到端多模态图内的模块选择独立的逻辑布局和物理秩放置。该系统支持两种不同的执行模式：

A. 核心抽象：边界通信器

为了在解耦模块布局的同时保留训练语义，系统引入了边界通信器。这些是附加在模块间边上的感知自动微分（autograd-aware）算子，负责处理不同分布式表示之间的张量转换。

• 前向传播：在目标模块的布局中实例化源激活。
• 后向传播：应用逆变换，将梯度路由回源模块的布局。
• 正确性：系统确保张量语义在独立网格间得到保留。通过严格定义每条边的逻辑批次分片转换（Equal-DP、Fan-in 和 Fan-out），防止了错误的映射。

B. 执行模式

1. 共置执行（Colocated Execution）：

   • 机制：模块共享相同的物理 GPU，但在不同的逻辑网格下解释它们（例如，具有高 DP 和低 TP 的编码器与具有低 DP 和高 TP 的 LLM 共享秩）。
   • 实现：通信器充当共享秩布局算子。它执行本地集合操作（例如 All-Gather 或区间选择），以在不跨秩集移动张量的情况下重新解释数据布局。
   • 调度：使用三阶段调度以避免编码器集合通信与交错式 LLM 流水线（1F1B）之间的冲突。
     • 阶段 1：编码器前向传播 + 共置前向转换。
     • 阶段 2：在转换后的张量上执行分离的 LLM 流水线（1F1B）。
     • 阶段 3：共置后向转换 + 编码器后向传播。
   • 优势：允许在 LLM 流水线阶段卸载编码器参数，从而提高内存效率。

2. 非共置执行（Non-Colocated Execution）：

   • 机制：模块占用不相交的秩集（例如，“编码器岛”和"LLM 岛”）。
   • 实现：使用确定性的边界领导者（具有规范坐标的秩）之间的显式点对点（P2P）通信。
     • 前向：源领导者发送给目标领导者；本地广播实例化数据。
     • 后向：梯度通过逆领导者路径路由回去。
   • 调度：使用感知图的 1F1B 调度器，在多阶段图中解析边，处理多个编码器流汇入 LLM 时的连接和分支。
   • 优势：防止编码器状态（参数、激活、优化器状态）与 LLM 在同一秩上竞争内存，从而允许 LLM 使用更高效的流水线布局。

3. 主要贡献

1. 异构并行运行时抽象：一个系统，允许单个 MLLM 图中的模块拥有独立的 TP/CP/PP/DP/EP 布局和物理秩放置，支持共置和非共置执行。
2. 边界通信器与调度：引入了感知自动微分的边界通信器以处理布局转换，并配合专用调度器（共置执行采用三阶段调度，非共置执行采用感知图的 1F1B 调度），以维持流水线效率。
3. 工作区评估：全面评估了异构性何时有益。研究发现，当内存余量充足但逻辑布局不同时，共置异构性是最优的；而当编码器状态限制 LLM 内存或流水线布局时，非共置异构性则变得更为优越。
4. 正确性与开源：通过逐步骤 FP32 检查和 LLaVA 风格的训练扫描，验证了异构布局与同质基线在损失收敛上的等价性。该系统作为 Megatron-LM 的开源扩展发布。

4. 实验结果

该系统在 NVIDIA H100 80GB GPU 上针对各种多模态工作负载（1B–12B 编码器，14B–120B LLM）和上下文长度（8K–128K）进行了评估。

• 共置异构性：
  • 在长上下文场景（例如 64K–128K，LLM 需要 CP 而编码器不需要）中，TFLOPS/GPU 最高提升 49.3%。
  • 在多种 1B/3B 编码器 + 14B/70B/120B LLM 配置中，通过允许编码器使用更高的 DP 和更低的 TP，吞吐量提升了 12.1%–21.6%。
• 非共置异构性：
  • 对于大编码器（6B–12B）+ 120B LLM 工作负载，聚合 Token 吞吐量最高提升 13.0%，TFLOPS/GPU 最高提升 9.6%。
  • 当编码器规模（例如 6B+）造成内存压力，从而在共置设置中限制 LLM 的流水线布局时，它成为首选工作点。
• 收敛性：验证了异构布局产生的损失曲线与同质基线完全相同，确认边界通信器正确保留了梯度语义。

5. 意义与主张

该论文主张，异构并行是多模态训练的必经演进，旨在摆脱僵化的模型级布局约束。通过暴露额外的自由度，该系统允许训练基础设施根据工作负载特征（编码器大小、上下文长度、Token 密度）动态调整工作点。

作者强调，这不仅仅是配置变更，而是一种根本性的运行时基础，它在保留端到端训练语义的同时，实现了：

• 共置执行：当模块可以共享硬件但需要不同的逻辑网格时，从固定的 GPU 预算中提取最大吞吐量。
• 非共置执行：当编码器状态否则会限制 LLM 时，解耦内存和调度压力。

这项工作作为 Megatron-LM 的开源扩展呈现，为未来的模块化 MLLM 框架和离散化训练系统提供了可复用的基础。该论文并未声称解决所有多模态挑战，而是专门解决了将异构模块强制纳入单一并行布局所导致的低效问题。