如果你曾经尝试过训练像 Llama 这样的大型语言模型 (LLM) 或像 ViT 这样的视觉巨型模型，你就知道其中的艰难: 单块 GPU 根本无法胜任。为了训练这些庞然大物，我们需要跨 GPU 集群进行分布式学习。

但问题在于: 仅仅拥有一个集群是不够的。你必须决定如何拆分模型。是拆分数据？拆分层？还是在层内部拆分张量？

历史上，这往往是在两种“两难”选择中做决定:

  1. 手动并行 (Manual Parallelism, MP) : 这要求你必须是硬件天才，需要手动调整每个张量在传输线路上的移动方式。
  2. 自动并行 (Automatic Parallelism, AP) : 试图由软件自动完成，但往往导致性能次优，因为它通常是将问题拆解成碎片来解决，而不是作为一个整体。

UniAP 登场。

在这篇文章中，我们将深入探讨一篇引人入胜的论文，该论文提出了一种利用混合整数二次规划 (MIQP) 来统一层间 (Inter-layer) 和层内 (Intra-layer) 并行的新方法。结果如何？该系统不仅能比现有的最先进方法更快地找到最佳并行策略，还能将训练吞吐量提高达 \(3.80\times\)。

问题所在: 并行策略的割裂

在介绍解决方案之前，我们需要了解分布式训练的现状。并行策略通常分为两类:

1. 层间并行 (Inter-Layer Parallelism)

这涉及在层与层“之间”有效地划分模型。

  • 流水线并行 (Pipeline Parallelism, PP) : 想象一条流水线。设备 A 计算第 1 层，将结果传递给设备 B，由设备 B 计算第 2 层。当 B 在工作时，A 开始处理下一批数据。

2. 层内并行 (Intra-Layer Parallelism)

这涉及在层“内部”拆分工作。

  • 数据并行 (Data Parallelism, DP) : 每个 GPU 都有一个模型副本，但处理不同的数据切片。它们在最后同步梯度。
  • 张量并行 (Tensor Parallelism, TP) : 单个巨大的矩阵乘法被拆分到多个 GPU 上进行。
  • 全分片数据并行 (Fully Sharded Data Parallelism, FSDP) : 一种内存高效的 DP 版本，其中参数、梯度和优化器状态被分片存储在各个工作节点上。

当前自动并行的局限性

现有的自动化工具 (如 Alpa 或 Galvatron) 通常采用分层 (hierarchical) 方法。它们可能先决定如何对模型进行流水线处理 (层间) ，然后再弄清楚每个阶段如何拆分张量 (层内) 。

这就像在确认能否付得起房租之前就决定住在哪个城市一样。通过按顺序解决这些问题，你很早就把自己锁定在一个次优解上。

UniAP (Unified Automatic Parallelism，统一自动并行) 采取了不同的方法。它同时考虑所有策略——PP、DP、TP 和 FSDP。

图 1 展示了 UniAP 与其他自动并行方法之间的区别。仅层间和仅层内 AP 方法从有限的集合中进行优化。分层 AP 方法按顺序优化。UniAP 拥有最大的联合优化策略空间。

如图 1 所示，UniAP 着眼于全局以找到真正的最优策略，而分层方法通常满足于“局部最优”。

UniAP 架构

那么，UniAP 如何在不耗费数年计算时间的情况下实现这种联合优化呢？研究人员将策略搜索视为一个数学优化问题。

工作流程如图 2 所示，包含三个主要阶段:

  1. 性能分析 (Profiling) : 收集有关模型和硬件的硬数据。
  2. 统一优化过程 (Unified Optimization Process, UOP) : 使用 MIQP 的核心引擎。
  3. 执行 (Execution) : 将数学结果转化为可运行的计划。

图 2. UniAP 流程图。系统获取硬件配置和模型架构，将其输入代价模型和 MIQP 求解器，并输出并行执行策略。

让我们拆解一下“UOP”黑盒内部发生的魔法。

1. 性能分析与代价建模

首先，UniAP 需要知道运行代价。它对硬件环境进行分析，以测量通信带宽 (P2P 和 All-Reduce 速度) 和计算能力。它还对模型进行分析，以了解每一层的内存占用和计算时间。

它使用两个代价模型:

  • 时间代价模型: 估算前向/后向传播时间和通信延迟。
  • 内存代价模型: 估算内存使用情况。

对于内存，特别是模型状态 (参数 + 梯度 + 优化器) ，UniAP 使用以下公式:

公式 1: 基于参数大小、张量并行大小和 FSDP 大小的模型状态内存代价。

这里，\(m_s\) 是内存代价，\(ps\) 是参数大小，\(ts\) 和 \(fs\) 分别代表张量并行和 FSDP 的大小。这个公式允许系统准确预测特定策略将消耗多少显存 (VRAM) 。

2. 混合整数二次规划 (MIQP)

这是论文的核心。研究人员将“如何快速训练该模型”的问题转化为一种称为混合整数二次规划的数学格式。

  • 混合整数 (Mixed Integer) : 某些决策是二元的 (0 或 1) 。例如，“第 5 层是否在 GPU 2 上？” (是/否) 。
  • 二次 (Quadratic) : 代价函数涉及变量相乘 (例如，将两个连接的层放在同一阶段会降低代价，这涉及二次项) 。

目标: 最小化 TPI

目标是最小化每次迭代时间 (Time Per Iteration, TPI) 。 在流水线系统 (如 GPipe) 中，速度取决于流水线的总延迟加上最慢阶段 (瓶颈) 所花费的时间。

图 3. GPipe 风格 PP 的时间代价分解。展示了跨阶段的前向 (fp) 和后向 (bp) 传播。

基于图 3 中的分解，目标函数如下所示:

公式 2: GPipe 中最小化 TPI 的目标函数。

该方程试图最小化处理时间 (\(p_i\)) 和通信时间 (\(o_j\)) 的总和，并根据微批次 (micro-batches) 的数量 (\(c\)) 进行加权。

约束条件

没有规则就没有优化。UniAP 施加了几个关键约束以确保解决方案有效。

A. 计算与通信约束 求解器计算每个阶段的代价。如果你将特定层分配给某个阶段，它们的计算代价会累加。如果层与层之间跨设备通信，那会增加通信代价。

公式 3: 每个阶段的计算代价约束。

公式 4: 阶段之间的通信代价约束。

B. 内存约束 这点至关重要。所选策略绝不能导致显存溢出 (OOM) 错误。任何设备上激活值和模型状态的总内存必须小于设备限制 \(m\)。

公式 5: 确保使用量不超过设备限制的内存约束。

C. 保序 (连续性) 约束 流水线并行不是随机的。如果你把第 1、2、3 层放在 GPU A 上，你通常不能把第 4 层放在 GPU B 上，然后又把第 5 层放回 GPU A。流水线阶段必须是层的连续集合。

图 4. 一个连续集合的例子。橙色节点代表有效的连续子图。

为了在数学上强制执行这一点，论文使用辅助变量 \(Z\) 来确保如果一个节点在一个阶段中，其依赖关系得到尊重。

公式 6: 确保每个计算阶段上的子图是连续的线性约束。

D. 放置与选择约束 最后，每一层必须被放置在某个地方，并且每一层必须只选择一种策略。

公式 7: 确保每一层都被分配到一个阶段的层放置约束。

公式 8: 确保每一层恰好选择一种并行策略的策略选择约束。

统一优化过程 (UOP)

定义好数学模型后，UniAP 运行一种算法 (论文中的算法 1) ，该算法:

  1. 枚举可能的流水线大小 (\(deg\)) 和微批次数量 (\(c\))。
  2. 将代价矩阵和约束条件输入求解器 (如 Gurobi) 。
  3. 求解器找到 TPI 的全局最小值。

可视化解决方案

该过程的输出是一组矩阵，准确指示每一层的去向及其使用的策略。

图 7. UOP 的候选解。显示了层到流水线阶段的映射以及层内策略 (DP/TP) 的选择。

如图 7 所示，矩阵 \(\mathbf{P}\) (放置) 和 \(\mathbf{S}\) (策略) 给出了执行的确定性蓝图。例如，第 0 层可能在第 1 阶段使用数据并行，而第 1 层使用数据并行和张量并行的混合。

矩阵表示在数学上如下所示:

公式 12: 放置 (P) 和策略 (S) 的矩阵表示。

实验与结果

研究人员在不同的硬件环境 (NVIDIA V100、A100 甚至 DCU) 下，使用多种模型 (BERT、T5、ViT、Swin、Llama) 将 UniAP 与顶级基线 (如 Galvatron 和 Alpa )以及手动方法 (如 Megatron-LM 和 DeepSpeed )进行了测试。

1. 吞吐量与优化速度

结果令人印象深刻。UniAP 不仅找到了更快的执行策略 (更高的吞吐量) ，而且找到策略的速度也快得多 (更低的优化时间) 。

表 1 & 2. 训练吞吐量和优化时间的比较。UniAP 在 NVIDIA 和 DCU 集群上均表现出色，优于或持平基线。

在表 1 中，注意 EnvC 上的 Llama-7B 。 UniAP 实现了 4.63 samples/s , 而 Galvatron 仅为 1.22。这是 3.80 倍的加速 。 为什么？UniAP 意识到在基于 PCIe 的 GPU 集群上，通信是瓶颈。它选择了 8 阶段流水线，并尽量减少张量并行，以减少跨设备通信。受限于其搜索空间的 Galvatron 选择了一种通信繁重的策略，阻塞了带宽。

此外，看看 BERT-Huge 的优化时间 。 UniAP 仅用 0.37 分钟就找到了解决方案，而 Alpa 花了超过 80 分钟。这就是 MIQP 求解器相对于复杂动态规划的效率优势。

2. 可扩展性

随着集群规模的扩大，UniAP 还能撑得住吗？

图 5. 可扩展性分析。(a) 吞吐量几乎呈线性增长。(b) 优化时间保持在可控范围内。

图 5(a) 显示吞吐量随节点数量线性增长。更重要的是, 图 5(b) 显示 UniAP 寻找策略所需的时间并没有呈指数级爆炸，这使其适用于大规模集群。

这种可扩展性即使在非 NVIDIA 硬件 (DCU) 上也成立，如表 4 所示:

表 4. DCU 集群上的可扩展性显示吞吐量呈线性增长。

3. 为什么联合优化很重要

研究人员进行了一项消融实验 , 以证明结合层间和层内并行实际上是必要的。他们将完整的 UniAP 与仅限于层间或仅限于层内策略的版本进行了比较。

图 6. 消融实验。限制策略空间会导致吞吐量显著降低或无可行解 (SOL X)。

图 6 中的结果非常鲜明。对于许多模型，限制搜索空间会导致性能极差或完全失败 (SOL \(\times\) 表示未找到解决方案，通常是由于 OOM) 。这验证了论文的核心前提: 必须进行联合优化。

案例研究: “EnvB”上的 BERT-Huge

为了具体说明，让我们看看 UniAP 在特定集群 (“EnvB”——2 个节点，节点间带宽较慢) 上为 BERT-Huge 找到的最佳策略。

图 9. BERT-Huge 的最佳并行策略。UniAP 智能地在节点间混合使用了 DP、TP 和 FSDP，以适应带宽限制。

如图 9 所示，UniAP 做了一些聪明的事情:

  • 它使用了 2 阶段流水线 (每个节点一个阶段) ，以最大限度地减少缓慢的节点间通信。
  • 在节点内部，它使用了张量并行 (TP) , 因为节点内带宽 (PCIe) 很快。
  • 它有选择地对特定层应用 FSDP (绿色框) ，以将所有内容放入内存中。

手动操作的工程师可能很难平衡所有这些因素，但数学模型自动解决了这个问题。

结论

分布式训练正在成为必需品，而非奢侈品。然而，配置这些系统的复杂性一直是一个主要障碍。

UniAP 代表了向前迈出的重要一步。通过将并行策略问题公式化为混合整数二次规划 (MIQP) 任务，它统一了层间和层内策略的搜索。主要结论如下:

  1. 统一 > 分层: 联合优化 PP、DP、TP 和 FSDP 比分步优化产生更好的结果。
  2. MIQP 速度快: 与传统搜索算法相比，数学求解器在导航这些巨大的搜索空间方面效率极高。
  3. 硬件感知: 无论是配备 NVLink 的 A100 集群还是基于 PCIe 的设置，UniAP 都能适应硬件的特定带宽和内存限制。

对于研究大型语言模型的学生和研究人员来说，像 UniAP 这样的工具为更易用、更高效和更自动化的模型训练铺平了道路。在白板上手动计算张量拆分的日子可能终于要结束了。