HexiScale: Accommodating Large Language Model Training over Heterogeneous Environment

作者/机构:
Ran Yan∗ (The Hong Kong University of Science and Technology), Youhe Jiang∗ (The Hong Kong University of Science and Technology), Xiaonan Nie (Peking University), Fangcheng Fu (Peking University), Bin Cui (Peking University), Binhang Yuan (The Hong Kong University of Science and Technology)

A1 主要贡献

核心问题：训练大型语言模型（LLM）是计算密集型任务，通常在具有同构高性能GPU的数据中心进行。然而，这种同构集群的部署成本高昂，限制了LLM的发展。因此，研究旨在探索一种替代方法，即将训练计算分布在异构GPU上，以实现更灵活高效的资源利用。

研究目标与挑战：在异构GPU上部署大规模LLM训练面临两大挑战：
1. 不同的GPU计算能力和内存容量：异构环境中GPU的浮点运算能力（FLOPS）和内存容量差异巨大。如果管理不当，高性能GPU可能未被充分利用，而低性能GPU则成为瓶颈。
2. 不同的GPU间网络带宽：GPU间的连接方式（如NVLink、PCIe、以太网）速度各异，导致通信时间不均衡，可能造成高速连接的设备等待低速连接，影响整体效率。

创新点与主要贡献：
为了应对上述挑战，本文提出了一个名为HexiScale的新型框架，用于协调异构GPU上的分布式LLM训练。其主要贡献如下：
1. 实现了支持非对称并行的HexiScale系统：该系统支持在数据并行、流水线并行和张量模型并行范围内对训练计算进行非对称划分，以灵活适应各种GPU的异构性和多样化的网络连接。这种设计能够将原始训练计算划分到合适的粒度，从而释放异构计算能力的全部潜力。
2. 提出了高效的调度算法：将异构GPU设备上的LLM训练计算分配问题形式化为一个约束优化问题。为高效求解该问题，提出了一种两阶段优化方法，该方法采用图划分算法来有效协调给定设备集的并行策略。第一阶段将可用GPU划分为多个组，每个组在第二阶段形成一个流水线。然后迭代应用该两阶段算法，以确定异构GPU集的最优并行策略。
3. 通过实验验证了HexiScale的有效性：在训练不同规模（7B、13B、30B）的Llama模型时，将HexiScale在异构环境下的系统效率与在同构数据中心运行的先进系统（Megatron、Galvatron、FSDP）进行了比较。结果表明，在总峰值FLOPS相同的情况下，HexiScale在异构环境中的模型FLOPs利用率（MFU）与同构环境下的SOTA系统相当，MFU百分比差距最低仅为0.3%，平均为3.5%。此外，HexiScale的性能也优于现有的异构训练系统（如Metis），MFU最高可达Metis的1.9倍。

A3 背景知识与系统设计

2.1 并行化LLM训练

并行训练策略。为了将LLM训练计算分布到数千个设备（通常是GPU）上，提出了三类主要的并行策略。

数据并行。数据并行（Data parallelism）【索引24，PyTorch distributed: experiences on accelerating data parallel training，2020，Proceedings of the VLDB Endowment】通过在设备间划分训练批次来分配计算，每个GPU都拥有一个模型副本用于前向和后向传播，梯度通过AllReduce操作进行同步。为了优化内存使用，可以将梯度、优化器状态和参数进一步分片到多个设备上，并在需要时通过额外的通信来收集。这类优化在Zero Redundancy Optimizer (ZeRO)【索引42，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】和Fully Sharded Data Parallel (FSDP)【索引59，PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel，2023，Proceedings of the VLDB Endowment】中得到了实现。

流水线并行。流水线并行（Pipeline parallelism）【索引10，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，Advances in neural information processing systems】、【索引31，PipeDream: generalized pipeline parallelism for DNN training，2019，Proceedings of the 27th ACM Symposium on Operating Systems Principles】将模型的计算按层划分为一系列阶段，每个GPU处理一个阶段。在前向传播期间，处理阶段i的GPU将激活值发送给处理阶段i+1的GPU；在后向传播期间，通信方向相反，以传递梯度。

张量模型并行。张量模型并行（Tensor model parallelism）【索引33，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】进一步将每个Transformer层划分到多个GPU上，其中权重矩阵按行或列进行分布。在前向传播中聚合层输出激活值和在后向传播中聚合相应梯度时，分别需要两次AllReduce操作。张量模型并行通信密集，但在GPU间连接速度快（例如，通过NVLink）时能有效并行化计算。

LLM训练的系统优化。各种系统优化【索引1，NeutronOrch: Rethinking SampleBased GNN Training under CPU-GPU Heterogeneous Environments，2024，Proceedings of the VLDB Endowment】、【索引5，Tensoropt: Exploring the tradeoffs in distributed dnn training with auto-parallelism，2021，IEEE Transactions on Parallel and Distributed Systems】、【索引13，Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks.，2018，ICML】、【索引14，Beyond data and model parallelism for deep neural networks，2019，Proceedings of Machine Learning and Systems】、【索引25，DAHA: Accelerating GNN Training with Data and Hardware Aware Execution Planning，2024，Proceedings of the VLDB Endowment】、【索引40，Accelerating Sampling and Aggregation Operations in GNN Frameworks with GPU Initiated Direct Storage Accesses，2024，Proceedings of the VLDB Endowment】已被提出，用于提高分布式LLM训练的吞吐量（每秒处理的token数），另一个广泛使用的衡量标准是模型FLOPs利用率（MFU），它衡量的是观测到的吞-吐量与系统利用100%峰值FLOPs时理论最大吞吐量的比率。

优化并行LLM训练的努力。为了优化并行LLM训练，已经做出了重大努力。零气泡并行（Zero bubble parallelism）【索引41，Zero Bubble (Almost) Pipeline Parallelism，2024，The Twelfth International Conference on Learning Representations】通过减少气泡开销来提高计算效率；引入了跨网格重分片机制（cross-mesh resharding mechanism）【索引61，On optimizing the communication of model parallelism，2023，Proceedings of Machine Learning and Systems】以最小化张量并行中的通信开销；此外，诸如梯度分桶、梯度累积和计算-通信重叠等系统优化已被集成到数据并行实现中【索引24，PyTorch distributed: experiences on accelerating data parallel training，2020，Proceedings of the VLDB Endowment】。

内存占用优化。为提高训练吞吐量，内存占用优化也十分必要。例如，激活重计算（activation recomputing）【索引6，Optimizing Large Model Training through Overlapped Activation Recomputation，2024，arXiv preprint arXiv:2406.08756】通过在后向传播期间重新计算所需的激活值而不是在前向计算后存储它，从而显著减少了内存占用。另一方面，将激活值卸载到CPU RAM或SSD【索引53，TBA: Faster Large Language Model Training Using SSD-Based Activation Offloading，2024，arXiv preprint arXiv:2408.10013】也可以通过自适应地将数据传输与计算重叠，从而在不影响性能的情况下有效减少GPU内存使用。

2.2 异构感知的LLM训练

针对异构资源的训练系统。学界已在构建专为异构计算资源定制的训练系统方面付出了巨大努力【索引29，Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism，2022，Proceedings of the VLDB Endowment】、【索引49，Unity: Accelerating {DNN} training through joint optimization of algebraic transformations and parallelization，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、【索引50，Improving Automatic Parallel Training via Balanced Memory Workload Optimization，2024，IEEE Transactions on Knowledge and Data Engineering】、【索引60，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】。例如，一些系统工作试图在异构和去中心化环境中普及和部署大型语言模型训练【索引27，Heterogeneity-aware distributed machine learning training via partial reduce，2021，Proceedings of the 2021 International Conference on Management of Data】、【索引55，Optimizing distributed training deployment in heterogeneous GPU clusters，2020，Proceedings of the 16th International Conference on emerging Networking EXperiments and Technologies】、【索引57，Decentralized training of foundation models in heterogeneous environments，2022，Advances in Neural Information Processing Systems】、【索引58，MiCS: near-linear scaling for training gigantic model on public cloud，2022，Proceedings of the VLDB Endowment】。SDPipe【索引28，Sdpipe: A semi-decentralized framework for heterogeneity-aware pipeline-parallel training，2023，Proceedings of the VLDB Endowment】实现了灵活的数据并行同步方案，以解决（半）异构环境中缓慢的数据并行通信问题；Whale【索引12，Whale: Efficient giant model training over heterogeneous {GPUs}，2022，2022 USENIX Annual Technical Conference (USENIX ATC 22)】提出了一种硬件感知的负载均衡算法，以加速在异构GPU上的训练。

与近期工作的比较。最近，Metis【索引48，Metis: Fast Automatic Distributed Training on Heterogeneous {GPUs}，2024，2024 USENIX Annual Technical Conference (USENIX ATC 24)】引入了一种新方法，可以自动为异构GPU上的分布式训练识别高效的并行方案。尽管这些系统可能是最相关的工作，但包括Metis在内的现有系统，由于在系统支持上存在局限，无法将原始训练计算划分到合适的粒度，因此未能充分利用异构计算能力的潜力（详见§5.4），这也导致缺乏高效的搜索算法来充分探索有效的调度空间。相比之下，HexiScale能够完全支持并行训练计算的非对称划分（见§3, §5），并拥有一个更高效、更有效的调度算法来找到接近最优的并行策略（见§4）。

3.1 案例研究：异构环境下的并行

场景设定。考虑在一个异构环境中训练一个Llama-2 (13B)模型，配置如下：机器A配备3块通过NVLINK连接的A800-80G GPU，机器内带宽为200 GB/s；机器B有3块通过PCIe连接的4090-24G GPU，带宽为32 GB/s；机器C配备2块通过PCIe连接的3090-24G GPU，带宽为16 GB/s。机器间通过1 GB/s的以太网链路互连。为了论证我们与Megatron相比的系统设计动机，我们模拟了一个全局批次大小为24，微批次大小为1的训练任务。两个系统都应用了激活重计算。

使用Megatron进行训练。Megatron只支持完全对称的训练计算划分。设$D_p, T_p, P_p$分别为数据、张量模型和流水线并行的并行度。在这8个GPU上考虑的并行策略包括：(i) 方案1：($D_p = 1, T_p = 1, P_p = 8$)；(ii) 方案2：($D_p = 1, T_p = 2, P_p = 4$)；(iii) 方案3：($D_p = 2, T_p = 1, P_p = 4$)。

方案1分析。方案1创建了一个有八个流水线阶段的流水线，由于高昂的气泡成本和不平衡的计算，效率低下。即使微批次大小设为1（这可能降低GPU核的效率），气泡时间仍占流水线执行时间的22%以上。此外，系统性能受到瓶颈机器B的拖累，其计算速度比机器A和C分别慢3.9倍和2.1倍。

方案2分析。方案2必须使用跨机器的张量模型并行。由于引入了每层约1.88秒的通信成本（在1 GB/s带宽和批次大小为24的情况下），性能受到严重影响。与使用NVLINK时的0.01秒通信成本相比，这是一个巨大的开销。

方案3分析。方案3是一个潜在的良好配置。我们通过最大化使用高机器内带宽进行数据并行通信的Transformer层数来微调此方案，如图1（上）所示。该策略创建了四个流水线阶段，每个阶段有10个层。然而，即使在这种设置下，Megatron也至少存在两个缺陷。首先，机器A中的两个GPU在流水线1上运行流水线并行，这浪费了高带宽的NVLINK，并引入了更高的流水线气泡开销。其次，机器A中的GPU处理的层数太少，导致计算能力的未充分利用。

Megatron的局限性总结。当底层异构环境与Megatron的完全对称策略不匹配时，系统不得不将计算能力强的GPU当作弱GPU使用，并且并行策略受到网络连接的限制，导致计算资源的利用不足。在这种情况下，先进的训练系统Megatron无法表现良好，估计的端到端训练迭代时间为41.52秒。

3.2 HexiScale中的非对称并行支持

核心设计变更。为了在异构设置下高效训练，我们实现了HexiScale，它支持完全非对称的并行和系统优化，其核心变化在于：

计算负载的非对称划分。对于每个可以分配不同批次大小的流水线，流水线并行通信组使用不同的张量模型并行度和相应可能不同的已分配Transformer层数进行初始化。每个流水线阶段选择一个固定的领导GPU，该GPU最小化到邻近阶段GPU的通信延迟，并初始化一个张量模型并行组。在前向传播期间，每个阶段的领导GPU将激活值发送到下一阶段的领导GPU。一旦下一阶段的领导GPU接收到激活值，它会在其张量并行组内广播此激活值以执行计算。在后向传播期间，对于激活值的梯度通信也应用相同的逻辑。

非对称梯度同步。不同流水线中的Transformer层的模型参数（及其对应的梯度）可能由于分配了不同的张量模型并行度而被分块成不同的大小。在这种情况下，普通的数据并行难以同步梯度。为了应对这一挑战，我们识别出最小的梯度大小，并将较大的梯度进一步划分为多个块，每个块的大小与识别出的最小梯度大小相匹配。然后，通过在数据并行组内的不同GPU子集中同步每个梯度块来执行数据并行通信，而不会增加通信开-销。

系统实现层面的优化。在系统实现方面，我们通过三种方式优化了我们的系统。首先，我们支持梯度累积和激活重计算，以降低数据并行通信成本和内存占用。其次，我们利用FlashAttention-2【索引7，FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning，2024，The Twelfth International Conference on Learning Representations】中的API来创建同时支持张量模型并行和flash attention的Transformer层。然后，我们用这些Transformer层实现了我们的非对称流水线并行。第三，我们通过注册自定义的FSDP【索引59，PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel，2023，Proceedings of the VLDB Endowment】通信钩子来支持非对称梯度同步逻辑。

3.3 案例研究：使用HexiScale提升性能

HexiScale如何提升效率。我们介绍HexiScale如何在之前的异构设置中提高训练效率。由于其灵活的设计，HexiScale可以更有效地构建两个流水线，如图1（下）所示。改进如下：

计算负载的完全划分。对于每个流水线，计算负载被完全划分如下：
* 从并行策略角度：我们通过在机器A上应用张量并行并分配更多的Transformer层，充分利用了其强大的计算能力和高机器内带宽。机器B仍然应用流水线并行，因为其内部连接速度不足以支持张量模型并行。
* 从层划分角度：机器B和机器C承担了较少的Transformer层，这解决了流水线阶段之间的计算不平衡问题。运行跨机器数据并行的层数被最小化。Megatron和HexiScale都受限于机器B上缓慢的数据并行通信。Megatron在机器B上为10个Transformer层进行通信，估计耗时9.90秒。而HexiScale只需要为5个Transformer层通信，耗时5.07秒，比Megatron快1.9倍。
* 从流水线效率角度：为平衡流水线之间的计算速度，分配了不同的批次大小。否则，端到端性能不会有改善，因为HexiScale仍将受限于较慢的流水线2。较快的流水线1必须等待流水线2进行数据并行通信。为了释放流水线1的效率，我们为流水线1分配了更大的批次大小。通过这种方式，我们平衡了每个流水线的运行时间，并提高了端到端的性能。在我们的案例研究中，两个流水线处理的批次大小相差40%，但运行时间仅相差7%。

运行非对称数据并行。对于机器A上流水线1的阶段0中的两个GPU，以及流水线1的阶段0和阶段1的GPU，尽管张量模型并行将参数划分为不同大小，但如§3.2所讨论的，支持在这些阶段之间运行数据并行通信。这种灵活的设计从两个方面提升了流水线1的性能。(i) 减少流水线阶段的数量可以降低流水线通信和气泡成本。(ii) 在机器A上采用张量模型并行可以增强流水线1的性能。

性能总结。总而言之，对于图1（下）所示的并行策略，HexiScale的端到端训练迭代时间估计为25.55秒，在这个假设的异构环境中比Megatron快1.6倍。

A2 方法细节

本节介绍我们的调度算法。（符号总结见表1。）

4.1 调度问题的形式化

问题定义。给定一组异构GPU，我们的目标是找到最优的异构并行执行策略，以最小化训练迭代时间。我们将调度问题形式化如下。设 $D = \{d_1, ..., d_N\}$ 是一组N个GPU设备，GPU设备内存限制记为 $m_d$。给定一个特定的GPU集，调度问题可以定义为找到最优的并行执行计划 $σ^∗$，该计划在内存消耗的约束下最小化训练迭代的执行时间：

$$\begin{aligned} \begin{aligned} \sigma^* = \arg \min_\sigma \quad & \text{\textsc{Comm-Cost}} (\sigma) + \text{\textsc{Comp-Cost}} (\sigma) \\ s.t. \quad & \text{\textsc{Mem-Cumsum}} (d) \leq m_d \quad \forall d \in \mathbf{D} \end{aligned} \end{aligned}$$

其中 Comm-Cost($σ$) 和 Comp-Cost($σ$) 表示给定并行执行计划 $σ$ 的一次迭代的通信和计算时间成本。Mem-Cumsum($d$) 表示设备 $d$ 的内存消耗。

问题的复杂性。一个并行执行计划 $σ$ 可以涉及任意数量的流水线，每个流水线具有不同的全局批次大小、微批次大小和并行策略。此外，一个特定的并行策略可以有不同的数据、流水线和张量模型并行度的配置。每个流水线阶段可以包含灵活数量的Transformer层。考虑到所有潜在配置的计算成本、通信成本和内存消耗，找到确切的最优并行策略是NP难问题，因为候选分配的规模呈指数级增长。

表1. 符号总结

解决方法概述。因此，在具有不同GPU容量和网络连接的异构集群中，通常只能使用基于启发式的调度算法来找到一个接近最优的并行执行计划。现有的调度算法通常无法找到高性能的并行策略，因为它们假设训练计算是对称划分的。例如，Alpa【索引60，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】中提出的算法，对网络连接和工作负载划分做了多个对称性假设，不能轻易地适应我们的场景。它们的搜索空间有限，且没有很好地考虑异构GPU之间的工作负载平衡。为了解决在异构集群中高效识别接近最优并行执行计划的挑战，我们设计了一个两阶段调度算法来寻找一个并行执行计划 $σ$，并迭代优化至一个接近最优的解 $σ^∗$。具体来说：
* 我们在§4.2中介绍第一阶段算法，该算法将设备集D划分为多个GPU组，每个组将用于创建一个流水线。
* 我们在§4.3中阐述第二阶段算法，该算法为每个流水线确定并行执行计划 $σ$。
* 我们迭代地重复两阶段算法，并将并行执行计划优化至 $σ^∗$，如§4.4所示。

4.2 调度算法的第一阶段

核心思想。第一阶段算法的关键思想是将GPU划分为多个组，每个组形成一个独立的流水线。然后，在这些流水线之间执行数据并行通信以同步梯度。

图划分方法。假设我们在当前迭代中将设备集D划分为 $D_{dp}$ 个GPU组，并最小化数据并行通信的带宽（我们将在§4.4中进一步讨论参数 $D_{dp}$ 和网络带宽分配）。我们首先将设备集D中的所有GPU组织成一个全局图 $G = (D, E)$，其中每个GPU $d \in D$ 是一个顶点，计算能力 $c_d$ 是顶点权重；对于任意 $d_1, d_2 \in D$，通信带宽 $β_{d,d'}$ 表示这两个GPU在集合E中的一条边。然后我们将全局图 $G$ 划分为一个全局分区 $P = \{D_1, D_2, ..., D_{D_{dp}}\}$，其中每个集合 $D_i \in P$ 包含用于第 $i$ 个流水线的高带宽GPU，并且对于任意 $i, j$，有 $D_i \cap D_j = \emptyset$。我们采用一种基于 $k$-way多级图划分算法【索引9，A Multi-Level Algorithm For Partitioning Graphs，1995，SC 95】的方法，该方法包括三个步骤：

步骤(i) - 粗化(Coarsen)。为了简化图划分，全局图被粗化成更小的图。直接将一个大的全局图（即包含许多顶点）划分为 $D_{dp}$ 个部分通常效率低下。相比之下，划分一个较小的粗化图则更有效。我们采用重边匹配（HEM）算法【索引21，A fast and high quality multilevel scheme for partitioning irregular graphs，1998，SIAM Journal on scientific Computing】。为了给数据并行分配低带宽，这种粗化操作意味着合并具有高带宽连接的GPU。如图2步骤(i)所示，粗化后的图的顶点数仅为全局图的一半。

步骤(ii) - 划分(Partition)。在这一步中，粗化后的图被进一步划分为 $D_{dp}$ 个GPU组，确保GPU组间的网络带宽最小化。我们采用递归二分法（recursive bisection method）【索引22，Multilevel algorithms for multi-constraint graph partitioning，1998，SC’98: Proceedings of the 1998 ACM/IEEE Conference on Supercomputing】，递归地对粗化图进行二分，直到获得 $D_{dp}$ 部分的划分。此步骤中的图划分解决了一个约束划分问题，即在保持严格平衡和精确划分为 $D_{dp}$ 个部分的约束下，最小化Cut目标函数【索引52，Towards efficient hierarchical designs by ratio cut partitioning，1989，1989 IEEE International Conference on Computer-Aided Design. Digest of Technical Papers】。Cut目标函数分两个层次定义：在层次一，任意两个集合 $D_i, D_j \in P$ 之间的Cut函数定义为连接它们的边权重（即带宽）之和。在层次二，全局分区P的Cut目标函数定义为任意两个集合 $D_i, D_j \in P$ 之间所有割的和。形式上，两级Cut函数可以定义为：

平衡约束。衡量全局分区 $P = \{D_1, D_2, ..., D_{D_{dp}}\}$ 平衡性的约束定义为顶点权重的最大和 $\max_{D_i \in P} \sum_{d \in D_i} c_d$ 与顶点权重的平均和 $\frac{\sum_{d \in D} c_d}{D_{dp}}$ 的比值。这个平衡因子总是大于或等于1。值越接近1，表示总顶点权重在GPU集 $D_i \in P$ 之间分布得越均匀。最大平衡因子被视为一个超参数。

步骤(iii) - 投影 & 步骤(iv) - 细化(Project & Refine)。划分粗化图并非最终目标——如图2步骤(iii)所示，为了找到全局图 $G$ 的划分，我们必须将结果投影回去，即应用步骤(i)的逆操作来恢复全局图 $G$ 中的 $D_{dp}$ 部分划分。为了有效考虑粗化节点内的信息，需要一个细化算法来提高划分质量并保持平衡；为此，我们在步骤(iv)中采用Kernighan-Lin算法【索引23，An efficient heuristic procedure for partitioning graphs，1970，The Bell system technical journal】来调整划分结果。

4.3 调度算法的第二阶段

核心思想。这一阶段的关键思想是根据第一阶段的图划分结果高效地生成流水线布局。得益于灵活的非对称数据并行设计，我们可以独立地为每个流水线确定并行策略。然而，搜索空间仍然很大，因为我们必须确定流水线和张量模型并行策略，以及流水线阶段的执行顺序。在完全异构的环境中，由于网络连接的异构性，仔细排列流水线阶段是必要的。

具体步骤。形式上，给定 $D_{dp}$ 组GPU，我们使用每个GPU集 $D_i \in P$ 中的GPU来为每个流水线找到并行执行计划 $σ_i$。为指定的GPU找到接近最优的布局涉及三个关键步骤：(i) 基于图划分对流水线阶段的GPU进行分组；(ii) 在每个GPU组内构建流水线阶段；以及 (iii) 在异构网络连接下确定流水线阶段的顺序。

步骤(i) - 为流水线阶段分组GPU。为了将具有高带宽连接的GPU分组，并为确定第 $i$ 个流水线的阶段顺序提供算法上的便利（稍后介绍），我们通过将集合 $D_i \in P$ 中的GPU进一步拆分为多个组来将具有高带宽连接的GPU分组。具体来说，我们首先将每个集合 $D_i \in P$ 组织成一个二级图 $G_i = (D_i, E_i), i = 1, ..., D_{dp}$，其中边集 $E_i$ 包含连接集合 $D_i$ 中GPU的通信带宽。接下来，我们将每个二级图 $G_i$ 划分为二级分区 $P_i = \{D_{i,1}, ..., D_{i,k_i}\}$，其中每个参数 $k_i$ 控制 $G_i$ 被划分成的部分数量。集合 $D_{i,j} \in P_i$ 包含用于构建流水线阶段的GPU，且对于任意 $j_1, j_2$，有 $D_{i,j_1} \cap D_{i,j_2} = \emptyset$。二级图使用与§4.2中讨论的相同的多级图划分方法进行划分。如图3（上）所示的示例，此步骤后，二级图中的GPU首先被划分为三个组。请注意，同一GPU集 $D_{i,j} \in P_i$ 内的GPU具有高带宽连接。在这些阶段中寻找阶段顺序只会带来微小的影响，因为流水线通信成本不是瓶颈。相反，当流水线阶段由不同的GPU集 $D_{i,j} \in P_i$ 生成时，排列这些阶段可以通过有效利用分配的低通信带宽来显著减少流水线通信开销。

后续步骤。在接下来的两个步骤中，我们首先在不考虑流水线阶段顺序的情况下，在每个GPU集 $D_{i,j}$ 内构建流水线阶段。然后，我们在由不同GPU集创建的流水线阶段之间搜索流水线阶段的顺序。

步骤(ii) - 构建流水线阶段。给定每个二级图 $G_i$ 的二级图划分结果，我们为第 $i$ 个流水线找到一个流水线布局，如下所示。每个GPU集 $D_{i,j} \in P_i$ 分别在每台机器内搜索其组内策略，这是通过使用附录A中定义的成本模型模拟不同的并行策略并选择局部最优的并行策略来完成的。如图3（中）所示，第一个GPU组构建了三个流水线阶段，而其他GPU组各自构建一个流水线阶段。

步骤(iii) - 通过贪心搜索找到流水线阶段顺序。我们将每个组内策略视为一个单独的顶点，并为第 $i$ 个流水线构建一个新的图 $G'_i$。第 $i$ 个流水线的阶段顺序因此通过一个top-τ贪心算法进行搜索。该算法在两个嵌套循环中运行。首先，它选择每个GPU组作为起始组。其次，对于每个具有τ个最高组间带宽的邻近GPU组，该算法递归地探索它们的邻近GPU组，直到生成一条流水线路径。如图3（下）所示，每个组内策略的阶段顺序没有改变（因为我们不排列它们），而第一个组内策略中的三个流水线阶段被放置在第二个组内策略的一个流水线阶段之后，以及最后一个组内-策略的一个流水线阶段之前。

4.4 迭代优化

优化过程。最后，我们介绍迭代优化过程——并行执行计划 $σ$ 从两个方面被迭代优化到最终的近优解 $σ^*$：

从第一阶段算法优化：第一阶段算法可以从两个方面进行优化。
首先，该算法通过在迭代中枚举参数 $D_{dp}$ 来将全局图划分为不同数量的流水线，以优化流水线的数量。

其次，该算法仔细地为数据和流水线并行通信分配网络带宽。我们的一个关键观察是，在异构环境中，数据并行通信和流水线执行时间都可能成为瓶颈。当流水线有许多阶段并处理大批量数据时，流水线执行时间占训练迭代时间的大部分。在这种情况下，最小化数据并行的带宽，并最大化流水线执行的带宽可以有效提高系统性能。系统性能得益于流水线执行通信开销的减少。另一方面，当流水线阶段少且处理小批量数据时，系统性能则得益于最大化数据并行的带宽，并最小化流水线执行的带宽。系统性能因数据并行通信开销的减少而得到提升。基于这一观察，我们实现了两种划分选项：(i) 最大化或 (ii) 最小化组间（即每个流水线的GPU组）的边权重（即带宽）。最大化组间边权重对应于最大化公式2中的Cut目标函数，这反过来导致为数据并行分配高通信带宽；相反，最小化边权重导致为数据并行分配低通信带宽。在每次迭代中，我们根据历史成本的移动平均值自适应地选择潜在的最优划分选项，以确保高效的系统性能。具体来说，我们在每次迭代结束时模拟数据并行通信和流水线执行成本，并更新历史平均成本。在下一次迭代中，使用这些历史信息来做出划分决策。通过这种设计，我们的第一阶段算法有效地分配了网络带宽，从而提升了系统性能。

从第二阶段算法优化：给定第一阶段获得的全局图划分，跨迭代地改变参数 $k_i$ 会导致每个组内策略的不同构建方式，这反过来决定了流水线阶段的配置及其顺序。通过微调 $k_i$，我们可以构建出高效的流水线。

性能评估。我们通过模拟来评估并行执行的性能：在每次迭代结束时，我们使用我们的成本模型模拟生成的并行执行计划的执行成本。当计划遇到内存不足问题时，执行成本被评估为无穷大。为提高准确性，我们进一步将网络延迟 ($α_{d,d'}$) 纳入模拟。随着集体通信操作数量的增加，系统性能因显著的链接成本而下降。网络延迟在异构环境中尤其关键，因为大量的微批次会导致NCCL操作数量增加。§5.3评估了模拟器的准确性，并显示所有情况下的模拟偏差均小于2%。

A4 实验

实验环境

• 模型：Llama-2 (7B), Llama-2 (13B), Llama (30B)。
• 基线框架：Megatron, Galvatron, FSDP。

• 硬件配置：

  • 同构设置：

    • 两台 8xA800 PCIe-80G 服务器（用于7B和13B模型测试），带或不带RDMA。
    • 四台 8xA800 PCIe-80G 服务器（用于30B模型测试）。

  • 异构设置（从Ucloud租用）：

    • 设置1 (用于7B, 13B模型)：一台 8x3080Ti (12G显存, 24GB/s内联)，一台 8x3090 (24G显存, 24GB/s内联)，三台 8x4090 (24G显存)。总FLOPS比两台8xA800 PCIe服务器高1.36%。
    • 设置2 (用于7B, 13B模型)：一台 8x3080Ti，一台 8x3090，一台 8x4090，一台 8xA800 NVLINK-80G (200GB/s内联)。总FLOPS比两台8xA800 PCIe服务器低1.59%。
    • 设置3 (用于30B模型)：一台 8x3090，两台 4x3090，四台 8x4090，一台 8xA800 NVLINK-80G。总FLOPS比四台8xA800 PCIe服务器低4.67%。

  • 网络：异构设置中的机器间连接为约0.7 GB/s的以太网。同构设置中测试了以太网和RDMA（10 GB/s）。

• 软件配置：

  • HexiScale基于FlashAttention-2和自定义FSDP通信钩子实现。

实验结果

5.1 端到端性能

• 同构环境性能：在同构设置下，HexiScale的性能与Megatron和Galvatron等高性能系统相当，尤其是在以太网连接下。FSDP由于是一维并行，在低速互联网络下表现不佳。（见图4和图5）

• 异构环境适应性：Megatron和Galvatron在异构环境中性能下降明显，甚至因内存不足（OOM）而无法运行（例如，在异构设置3中训练30B模型）。原因是它们的对称并行策略导致了严重的负载不均衡和资源浪费。Galvatron虽比Megatron灵活，但其并行策略对异构集群仍非最优。相比之下，HexiScale在异构环境中表现出色，其最大MFU比Galvatron高出2.5倍，平均提升2.1倍。（见图5）

• 异构 vs. 同构性能对比：为了公平比较，将HexiScale（异构+以太网）与基线系统（同构+以太网）进行对比。尽管异构集群面临GPU内存不均衡和更频繁的跨机通信等挑战，HexiScale的MFU与同构场景相比，差距最小仅为0.3%，平均差距为3.5%。这表明HexiScale通过非对称设计和有效调度，能高效利用碎片化的GPU资源。（见图4和图5）

5.2 消融研究

• 系统设计分解：

  • 非对称并行支持：禁用非对称并行支持后，系统性能在训练Llama(30B)时最多下降23%，平均下降15%。这是因为对称策略无法充分利用不同硬件的特性（如图6）。
  • 梯度累积(GA)：禁用梯度累积后，系统性能最多下降15%，平均下降12%。这是因为异构GPU内存有限，批次大小受限，若无GA会导致频繁的数据并行通信，降低性能（如图6）。

• 训练迭代延迟分解：与Galvatron相比，HexiScale显著减少了通信开销和流水线气泡。Galvatron的对称流水线配置导致了高并行度，增加了通信和气泡开销，并因计算不平衡产生额外的“拖后腿”气泡。HexiScale通过减少流水线阶段数和精心平衡计算来解决这些问题。（见图7）

5.3 调度算法评估

• 调度结果案例：表2展示了调度算法为Llama(30B)在异构设置3中找到的并行策略。该策略包含4个流水线，每个流水线内部根据硬件特性应用了不同的局部最优并行策略（如在A800上使用高TP度）和层分配，以平衡计算负载。

• 模拟准确性：表3显示，模拟的MFU结果与实际运行结果非常接近，偏差小于2%，证明了算法能够准确地搜索有效的并行执行计划。

• 图划分有效性：与随机图划分相比，本文提出的图划分算法在搜索Llama-2(7B)和Llama(30B)的并行计划时，能收敛到更高的估计MFU，分别高出约8%和23%。此外，该算法能持续迭代改进，而随机划分则容易早早收敛且性能较差。（见图8）

• 算法可扩展性：算法在扩展到大规模集群时表现出良好的效率。在12核CPU上对64到320个GPU进行50次迭代搜索，运行时间保持在两分钟以内，开销远小于LLM训练时间。同时，模拟的MFU表明算法在大型集群上性能稳定。（见图9）

表2. HexiScale在异构设置3中训练Llama (30B)时发现的并行策略。共有四个流水线(DP=4)，具有不同的流水线布局。

表3. 不同实验设置下真实与模拟MFU的比较。

5.4 案例研究

• 与Metis的比较：在异构设置3中，与另一SOTA异构训练系统Metis进行比较。Metis的调度算法倾向于在每个流水线阶段使用高数据并行度，这导致其流水线气泡开销显著增加。实验结果显示，Metis的MFU为17.1%，比HexiScale低1.6倍。（见图10左）

• 延迟分解对比：图11显示，与Metis相比，HexiScale的通信开销略高，但流水线气泡时间显著减少。这是因为HexiScale通过减少流水线阶段数来有效缓解气泡问题，从而获得更好的整体性能。

• 大规模模拟：在一个包含240个异构GPU的大型集群上，对Llama(30B)和Llama-2(70B)模型进行了性能模拟。结果显示，HexiScale的性能分别比Metis高1.6倍和1.9倍，比Galvatron高2.7倍和3.3倍。此外，Metis和Galvatron的搜索算法在大型集群上效率低下，需要数小时才能完成。（见图10中、右）

A5 结论

本文介绍了HexiScale，一个利用异构GPU提升LLM训练灵活性和效率的新型系统。通过支持数据、流水线和张量模型并行的非对称划分，HexiScale有效地利用了多样化的GPU能力。我们通过一个分层图划分算法来优化这些非对称计算。实验研究表明，在训练7B到30B参数的模型时，HexiScale在异构GPU上实现了与使用同构高性能GPU的SOTA系统相当的吞吐量。HexiScale的性能也优于SOTA的异构训练系统。这些结果凸显了HexiScale通过利用异构GPU使LLM训练更易于获取和更具成本效益的潜力。

A6 附录

A 成本建模

分步建模。本节中，我们逐步对通信成本（Comm-Cost）、计算成本（Comp-Cost）和内存消耗（Mem-Cumsum）进行建模。首先，我们为每个Transformer层建模成本，然后为整个模型建模端到端成本，如下所示：

逐层成本建模

• 张量模型并行通信成本。假设启用了激活重计算。一个Transformer层在一组GPU $d_{i,j}$ 上运行，一个微批次的张量模型并行通信成本可以估计为：

  $$12 \cdot \max _{d \in \mathbf{d}_{i, j}} \sum_{d^{\prime} \in \mathbf{d}_{i, j}-\{d\}}\left(\alpha_{d, d^{\prime}}+\frac{B_{m b} S H B_{\text {type }}}{\left|\mathbf{d}_{i, j}\right| \beta_{d, d^{\prime}}}\right)$$

• 数据并行通信成本。一个Transformer层的数据并行通信成本可以估计为：

  $$ \text{COMM-DP-LAYER} \left( \mathbf{D}_{\mathrm{dp}}^{k} \right) = 2 \cdot \max_{d \in \mathbf{D}_{\mathrm{dp}}^{k}} \sum_{d^{\prime} \in \mathbf{D}_{\mathrm{dp}}^{k} - \{d\}} \left( \alpha_{d, d^{\prime}} + \frac{12 H^{2} B_{\text{type}}}{\left| \mathbf{D}_{\mathrm{dp}}^{k} \right| \beta_{d, d^{\prime}}} \right) $$

• 流水线并行通信成本。请注意，流水线并行通信仅在两个层位于不同阶段时发生，例如，第 $i$ 个流水线中的第 $j$ 阶段和第 $j+1$ 阶段。这可以视为两个步骤：在前向传播中，第 $j$ 阶段发送到第 $j+1$ 阶段，然后第 $j+1$ 阶段将信息广播到该阶段的每个GPU（或者在后向传播中，第 $j$ 阶段从第 $j+1$ 阶段接收信息然后广播）。为方便起见，定义 Comm-PP-Hop($d_{i,j}, d_{i,j+1}$) = 0。一个微批次的通信成本可以估计为：

  

• 一个微批次的计算成本。假设张量模型并行总是在同一种类型的设备 $d$ 上运行，并且启用了激活重计算。计算成本可以估计为：

  

每种并行策略的成本建模

• 数据并行成本。不同的流水线阶段同时同步梯度。数据并行成本受最慢的流水线阶段限制，可以估计为：

  $$\text{COMM-DP} = \max_{i,j} \left[ \sum_{k=s_{i,j}}^{l_{i,j}} \text{COMM-DP-LAYER} \left( \mathrm{D}_{\text{dp}}^k \right) \right]$$

• 流水线和张量模型并行成本。对于第 $i$ 个流水线的执行，其成本包括每个阶段（由 $j$ 索引）的计算和通信成本：

  

  请注意，最慢的阶段限制了流水线并行的阶段。因此，我们将流水线和张量模型并行成本公式化如下：

  $$\begin{aligned} \begin{aligned} & \sum_{j=1}^{D_{p p}^{i}}\left(\operatorname{STAGE}\left(\mathbf{d}_{i, j}\right)+\operatorname{CoMM-PP-HOP}\left(\mathbf{d}_{i, j}, \mathbf{d}_{i, j+1}\right)\right) \\ & +\left(n_{m b}-1\right) \cdot \max _{j=2, \ldots, D_{p p}^{i}}\left(\operatorname{STAGE}\left(\mathbf{d}_{i, j}\right)\right. \\ & \left.+\operatorname{CoMM-PP-HOP}\left(\mathbf{d}_{i, j}, \mathbf{d}_{i, j+1}\right)\right) \end{aligned} \end{aligned}$$

端到端时间建模：一次迭代的时间由最慢的流水线和数据并行成本决定，可以估计如下：

$$\begin{aligned} \begin{aligned} & \text{\textsc{Comm-Cost}} (\sigma) + \text{\textsc{Comp-Cost}} (\sigma) \\ & \quad = \max_{i=1, \dots, D_{dp}} \text{\textsc{Pipeline-Time}} (i) + \text{\textsc{Comm-DP}} \end{aligned} \end{aligned}$$

内存成本建模：假设应用了完全激活重计算和朴素的数据并行。参数和激活的内存成本可以估计如下：

$$ \text{MEM-CUMSUM}(\sigma) = \frac{48 H^2 B_{type}}{|\mathbf{d}_{i,j}|} + B_{mb} SHB_{type} $$

B 实验设置细节

大规模集群模拟。我们在一个由240个异构GPU组成的大型集群上模拟了HexiScale、Metis和Galvatron的性能，具体配置如下：

表4. 模拟的大型集群中的GPU组成。

C 并行策略细节

异构设置3的策略。Metis和Galvatron在异构设置3中的并行策略如下：

C.1 Metis的策略

策略偏好。Metis倾向于使用机器内数据并行，因此构建了一个包含8个阶段的单一流水线，如下：

表5. 异构设置3下Metis的并行策略。

表6. MILP和HexiScale在不同模型和异构设置下的MFU比较。

C.2 Galvatron的策略

策略调优。Galvatron的策略被微调为 $(D_{dp}, P_p, T_p) = (4, 2, 7)$。如果数据并行度更高，内存不足问题持续存在。如果数据并行度更低，系统性能会因显著的流水线气泡或通信开销而受损。张量模型并行度被微调为2，否则，对于更高的张量模型并行度，3090会产生显著的通信开销，并损害整体系统性能。

D 算法最优性

问题性质。由于这个复杂问题的NP难特性，确定一个上界尤其具有挑战性，因为它需要一个可证明的最优解或在特定约束下的启发式保证。

与MILP的比较。为了评估我们调度算法的最优性，我们将其与混合整数线性规划（MILP）方法进行比较，这是一种在资源分配问题中广泛采用的方法。基于ILP/MILP的求解器已被集成到Alpa【索引60，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、FlexSP【索引51，Data-Centric and Heterogeneity-Adaptive Sequence Parallelism for Efficient LLM Training，2024，arXiv preprint arXiv:2412.01523】和Helix【索引26，Helix: Distributed Serving of Large Language Models via Max-Flow on Heterogeneous GPUs，2024，arXiv preprint arXiv:2406.01566】等著名系统中，证明了它们在LLM训练和推理的资源分配和并行策略确定方面的可靠性和有效性。

性能评估。我们同时运行了MILP算法和我们的算法，并比较了性能差距。MILP算法保证了最优结果，但其运行时间呈指数增长。

结果分析。在表6中，我们模拟了§5.1中出现的异构设置下的系统性能。如表所示，HexiScale的调度算法实现了与MILP相当的性能，MFU的性能差距小于2%。此外，HexiScale使用的基于图划分的调度算法在处理复杂和大型环境时具有出色的可扩展性，在搜索时间上显著优于MILP（HexiScale在几分钟内完成搜索，而MILP需要数小时）。总之，我们的算法提供了接近最优解的性能，同时表现出卓越的效率。

E 网络变化的影响

算法适应性。HexiScale通过提出的调度算法，旨在适应不同的网络条件（以及异构的计算能力）。

策略调整。当网络发生变化时，HexiScale将灵活地配置并行策略。直观上，高带宽应该分配给通信密集的并行策略，而低带宽则应避免或分配给通信量较小的并行策略。

具体示例。为了提供一个具体的例子，我们通过改变异构设置1（40个异构GPU）中Llama-2 (7B)模型的机器间带宽来模拟系统性能，如我们论文§5.1所述。模拟结果显示，HexiScale在0.5 GB/s的机器间带宽下实现了30.1%的MFU，在0.7 GB/s下实现了32.3%的MFU，在5 GB/s下实现了36.4%的MFU。

低带宽场景。当机器间带宽从1 GB/s降至0.5 GB/s时，并行策略和网络分配保持不变。流水线并行仍然使用较慢的机器间带宽进行通信，而数据并行则分配给较高的机器内带宽。由于流水线通信量相对较小，降低的网络带宽仅引入了少量额外的开销，这反映在MFU上。

高带宽场景。当机器间带宽从1 GB/s增加到5 GB/s时，调度算法将重新分配机器间带宽给数据并行，因为数据并行的通信量要大得多。这一调整通过减少与数据并行相关的通信开销，显著提升了系统性能。

总结。总之，我们的算法通过调整并行策略，并根据通信量为每个策略分配合适的网络链接来适应网络变化，从而优化计算利用率。