笔记
智能计算软件优化
分布式MLSYS

分布式MLSYS

1. 引言：大模型发展的挑战与趋势

本章节旨在介绍当前大模型发展的宏观背景、驱动因素以及所面临的系统性挑战，特别是计算硬件发展与模型规模增长之间的矛盾，从而引出分布式训练系统的重要性。

1.1 Scaling Law：模型性能与规模的关系

  • 核心定律：Scaling Law 指出，在其他条件不变的情况下，模型越大、数据越多、计算越久，模型的表现就越好。
    • 这意味着通过简单地增加模型的参数量、训练数据集的大小以及训练时间，可以系统性地提升神经网络模型的性能。
  • 重要参考文献：这一观察结果在神经网络语言模型领域得到了广泛验证，其中一篇里程碑式的论文是 Kaplan 等人发表的《Scaling Laws for Neural Language Models》（https://arxiv.org/pdf/2001.08361.pdf）。
  • 深远影响：Scaling Law 是推动当前大模型浪潮（尤其是大型语言模型LLM）发展的重要理论基础。它为研究人员和工程师提供了明确的优化方向，即通过规模化来实现性能突破。

1.2 LLM维度增长：上下文大小的扩展

  • 关键维度：除了模型参数量的增长，增加模型的上下文大小（Context Size） 是LLM性能提升的另一个关键维度。
    • 上下文大小决定了模型在进行预测或生成文本时能够考虑的历史信息量。更大的上下文意味着模型可以理解更长距离的依赖关系，处理更复杂的语境，从而在长文本理解、对话连贯性、代码生成等任务中表现更优。
  • 重要参考文献：Brown 等人发表的《Language Models are Few-Shot Learners》（https://arxiv.org/pdf/2005.14165.pdf）展示了大型语言模型在具备足够上下文时所展现出的强大的“少样本学习”（Few-Shot Learning）能力。

1.3 计算硬件发展与模型规模增长的矛盾

  • 现状：当前大模型的模型规模增长速度远远超过了计算硬件性能的提升速度，二者之间存在显著的“剪刀差”。
    • GPU性能增长：GPU 的 FLOPS（每秒浮点运算次数）性能大约每 2.5年翻一番。
    • LLM规模增长：然而，大型语言模型的规模（参数量）却以每年约10倍的速度增长。
  • 影响：这种巨大的差距意味着，即使拥有最先进的单个计算设备，也难以满足训练和部署超大规模模型的需求。
  • 重要参考文献：Hobbhahn 和 Besiroglu 在《Trends in GPU Price-Performance》（https://epochai.org/blog/trends-in-gpu-price-performance）中对这一趋势进行了详细分析。此外，Smith 等人关于训练 Megatron-Turing NLG 530B 的论文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》（https://arxiv.org/pdf/2201.11990.pdf）也侧面反映了应对这种规模挑战所需的系统级努力。
  • 结论：为了弥合硬件能力与模型需求之间的鸿沟，分布式系统成为了承载和加速大模型计算的必然选择。

1.4 大模型参数规模演变图

该部分原始材料中包含一张展示大模型参数规模随时间演变趋势的图表。以下是对该图表的详细描述和解读：

  • 图表内容描述：
    • 图表横轴代表年份，从2014年（经典机器学习模型）一直延伸到2024年（预留发展空间）。
    • 图表纵轴代表模型参数规模（Scale of parameters），从数千万到数万亿。
    • 图表上标记了多个里程碑式的模型，展示了它们在不同年份的参数量。
    • 图表下方还列举了为了支持更大模型训练而引入的各种分布式并行技术。
  • 关键信息提取：
    • 2014年及以前：标记为“Classic ML models”，参数规模相对较小，通常在百万级。
    • 2018年：ELMo (94M) - 参数量达到亿级。
    • 2019年：BERT-L (340M), GPT-2 (1.5B) - 参数量进入十亿级。
    • 2020年：Megatron-LM (8.3B), Turing-NLG (17.2B), GPT-3 (175B) - 参数量迅速增长，GPT-3达到千亿级。
    • 2021年：Switch-C (1.6T), GShard (400B) - Switch-C首次将参数量推向万亿级别（稀疏模型）。
    • 2022年：Megatron-Turing (530B) - 继续在千亿级规模发展。
    • 2023年：图表指示参数规模继续增长。
  • 分布式技术的重要性：
    • 随着模型参数的爆炸式增长，单设备已无法满足存储和计算需求。图表下方列出的分布式并行技术正是应对这一挑战的关键。
    • Data Parallel (数据并行)：最早且最常用的并行策略，通过在不同设备上复制模型并分配不同批次的数据来加速训练。
    • Pipeline Parallel (管道并行)：将模型的层（layers）分布到不同的设备上，形成一个处理流水线，以减少单设备的内存占用。
    • Tensor Parallel (张量并行)：将模型内部的张量（如权重矩阵）切分到不同设备上，实现层内的并行计算。
    • Sequence Parallel (序列并行)：在序列维度上对输入进行切分，以减少激活内存的消耗。
    • Expert Parallel (专家并行)：主要用于稀疏激活模型（如MoE），将不同的“专家”模型分配到不同设备上，只激活部分专家进行计算。
  • 总结：这张图生动地展示了**“Scaling Law”**的实践结果，以及在模型规模不断攀升的背景下，分布式训练系统和多样化的并行策略从“需要”变为“必须”的发展历程。这些并行策略的组合使用使得训练万亿参数规模的大模型成为可能。

2. 通讯原语与通讯机制

在分布式机器学习训练中，数据需要在多个计算设备（如GPU）之间高效地传输和同步。为了实现这一目标，需要依赖一系列底层的通讯原语和机制。本节将详细介绍这些核心概念。

2.1 集合通讯 (Collective Communication)

2.1.1 定义与作用
  • 定义：集合通讯（Collective Communication）指的是一组参与进程（或设备）之间协作完成的通信操作。它们需要共同完成某种“整体性”的数据传递或计算任务。
  • 作用：在大规模并行计算中（如多GPU训练），许多操作具有全局性质，例如聚合所有设备的梯度或分发相同的模型参数。这些操作无法简单地通过点对点通信（即两个设备之间的直接通信）高效完成，因此需要更高级别的集合通讯操作来协调多个设备。集合通讯提供了一种抽象，使得开发者无需关心复杂的点对点通信细节。
2.1.2 通讯抽象层与实现层

分布式深度学习框架通常采用分层设计，将复杂的通信操作进行抽象：

  • 上层通信抽象层：
    • 提供给用户友好的API，如 torch.distributed (PyTorch), Horovod, DeepSpeed。
    • 这些框架封装了底层的通信细节，允许开发者以更高级别的语义（例如“同步梯度”、“广播模型”）来组织分布式训练逻辑。
  • 下层通信实现层：
    • 负责具体的网络数据传输，提供高性能的通信后端。
    • 常见的实现包括：
      • NCCL (NVIDIA Collective Communications Library)：NVIDIA专门为GPU优化的高性能集合通信库，广泛用于多GPU和多节点GPU集群。
      • MPI (Message Passing Interface)：通用的并行计算通信协议，可在CPU和GPU集群上使用。
      • Gloo：Facebook开源的CPU通信库，也可以在GPU上运行。
      • RCCL (ROCm Collective Communication Library)：AMD为自家GPU平台提供的集合通信库。
      • UCX (Unified Communication X)：一个高性能的通用通信框架，支持多种网络硬件和传输协议。
      • SHARP (Scalable Hierarchical Aggregation and Reduction Protocol)：Mellanox InfiniBand网络提供的硬件加速集合通信技术。
      • RDMA (Remote Direct Memory Access)：一种允许直接访问远程内存而无需CPU参与的技术，能够显著降低通信延迟和CPU开销。
      • TCP (Transmission Control Protocol)：标准的网络传输协议，但通常用于低带宽或广域网场景，性能不如专门的并行通信协议。
2.1.3 常见的集合通讯操作

以下是分布式训练中常用的几种集合通讯操作的列表，它们将在2.3节中详细讲解：

  • AllReduce
  • Broadcast
  • Reduce
  • AllGather
  • ReduceScatter

2.2 基于点对点通讯实现的通讯原语

除了上述直接由底层库提供的集合通讯原语，许多更复杂的通信模式可以通过基本的点对点通讯（如 ncclSend() 和 ncclRecv()）进行构建和实现。这些原语描述了数据如何在特定发送方和接收方之间一对一地传输。基于点对点通讯，可以构建出：

  • Scatter (One-to-all)：将一个设备的数据分散到所有其他设备。
  • Gather (All-to-one)：将所有设备的数据收集到一个设备。
  • All-to-all：每个设备都向其他所有设备发送数据，并从其他所有设备接收数据。

2.3 各类通讯原语详解

2.3.1 AllReduce
  • 定义与功能：
    • Reduce 指的是在某个或某些维度上对数据进行归约处理（如求和 sum、求最小值 min、求最大值 max、求平均 avg 等操作）。
    • AllReduce 则是在完成归约处理后，将最终的归约结果分发（复制）给所有参与集合通讯的GPU（或进程）。
    • 典型应用：在数据并行训练中，每个GPU会计算各自子批次数据的梯度，然后使用AllReduce操作将所有GPU计算出的梯度进行求和（或求平均），并将最终的平均梯度同步到所有GPU上，以便每个GPU都能用相同的梯度更新模型。
  • 实现优化：并行传输，管道传输：
    • 为了提高效率，AllReduce的实现通常会采用并行传输（例如同时在多个链路上发送数据）和管道传输（将数据分成小块，在前一块数据传输的同时，后一块数据进行计算或准备发送，形成流水线效应）等技术。
  • Ring-based AllReduce: Reduce-Scatter + Allgather：
    • 环形AllReduce是一种非常高效且常用的AllReduce实现算法，尤其适用于大规模集群。
    • 它将一个完整的AllReduce操作分解为两个阶段：
      1. Reduce-Scatter：每个设备将自己的数据与环上的下一个设备的数据进行归约，并将归约结果的一部分（slice）发送给下一个设备。这个过程在环上进行N-1次，直到每个设备都获得了所有设备数据的一个归约后的“片段”（slice）。最终，每个设备都拥有一个完整的归约结果的不同片段，且这些片段已经包含了所有设备对应部分的数据。
      2. Allgather：在Reduce-Scatter完成后，每个设备都拥有归约结果的一个片段。Allgather阶段则通过在环上传输这些片段，使得每个设备最终都能收集到所有片段，从而得到完整的归约结果。
    • 优点：环形AllReduce能够有效地利用网络带宽，并通过流水线操作隐藏通信延迟，使得通信流量在所有设备上相对均衡。
2.3.2 Broadcast
  • 定义与功能：
    • Broadcast 操作是将一个设备（源设备）上的数据完全复制并发送到所有其他参与集合通讯的设备上。
    • 典型应用：在分布式训练中，通常用于将主设备（或某个特定设备）上的模型参数、超参数或配置信息分发给所有worker设备。
  • 实现优化：并行传输，管道传输：
    • 与AllReduce类似，Broadcast的实现也可以通过并行传输和管道传输来优化效率。例如，可以使用树形结构进行广播，以减少传播时间。
2.3.3 Reduce
  • 定义与功能：
    • Reduce 操作是将所有设备上的数据进行归约，并将最终的归约结果只存储到其中一个指定的设备上。
    • 与AllReduce的区别：“没有All前缀”意味着归约结果只在一个位置（通常是根节点）可见，而不会像AllReduce那样分发到所有设备。
    • 典型应用：收集所有worker设备的日志信息、统计量或汇总某些特定的中间结果到主设备进行统一处理。
2.3.4 Scatter (One-to-all)
  • 定义与功能：
    • Scatter 操作是将一个设备上的数据分割成多个不重叠的片段（chunks），然后将这些不同的片段分发给不同的参与设备。
    • 典型应用：将一个大的数据集或一个大型批次的数据分割成多个小批次，分发给不同的worker设备进行独立处理。
  • 与Broadcast的区别：
    • Broadcast 发送给所有设备的数据是完全相同的副本。
    • Scatter 发送给所有设备的数据是不同且不重叠的片段。
2.3.5 ReduceScatter
  • 定义与功能：
    • ReduceScatter 操作结合了Reduce和Scatter的功能。它首先将所有设备的数据进行归约，同时在归约过程中，将归约结果的不同片段直接分散到不同的设备上。
    • 换句话说，每台设备都有自己的数据。在将每个设备的数据分割分发的过程当中同时完成归约操作。
    • 典型应用：在环形AllReduce算法的第一阶段中用到，每个设备将自身数据与环上下一个设备的数据归约，并将归约结果的一个片段发送出去。
2.3.6 Gather (All-to-one)
  • 定义与功能：
    • Gather 操作是将所有参与设备上的数据收集到其中一个指定的设备上。这些数据通常是来自不同设备的，收集后会在目标设备上按照特定顺序组合起来。
    • 可以形象地理解为将分散在不同设备上的数据“行变列”或“拼接到一起”。
    • 典型应用：在模型训练中，每个设备可能处理了不同的数据片段并生成了对应的结果。Gather可以将这些结果收集到一起进行后续处理，例如在主设备上对所有worker的输出进行统一评估。
  • 与Scatter的逆操作关系：
    • Gather和Scatter互为逆操作。Scatter将一个完整的数据源分散到多个目标，而Gather则将多个分散的数据源收集回一个完整的目的地。
2.3.7 AllGather
  • 定义与功能：
    • AllGather 操作是在Gather收集数据的基础上，将最终收集到的完整结果（所有设备的数据组合而成）复制并分发给所有参与集合通讯的设备。
    • 典型应用：在需要所有设备都拥有完整信息（例如所有设备的输出特征，或在张量并行中需要组合完整的张量）的场景。
2.3.8 All-to-All
  • 定义与功能：
    • All-to-All 是一种更通用的通信模式，可以理解为Scatter和Gather同时发生。
    • 每个设备将其本地数据分割成多个块，并将每个块发送给不同的目标设备。同时，每个设备也从所有其他设备接收数据块，并组合成其最终的输出数据。
    • 形象地说，如果将所有设备的数据看作一个矩阵，All-to-All操作相当于对这个矩阵进行“行列互换”（转置），即每一列是某个设备Scatter的结果，每一行是某个设备Gather的结果。
    • 典型应用：在张量并行和专家混合模型（MoE）中，经常需要进行All-to-All通信，以实现数据或模型参数的复杂重排和交换。例如，MoE中将tokens路由到不同的专家，每个专家处理一部分tokens，然后将处理结果再路由回来。

3.1 数据并行 (Data Parallelism, DP)

3.1.1 定义与基本原理

数据并行 (Data Parallelism, DP) 是一种分布式训练策略，它通过将大规模训练数据的批次（batch）切分为多个子批次（mini-batch），然后将这些子批次分发到不同的计算设备（如GPU）上并行处理。每个设备都拥有模型的一个完整副本，独立地对分配到的数据子集进行前向传播和反向传播，计算出各自的梯度。最后，这些设备上的梯度会被聚合（通常是求平均），并用于更新所有设备上的模型副本。这个过程与在单个设备上对原始大批次数据进行梯度计算是等价的，从而加速了训练过程。

3.1.2 DP的计算流程

数据并行的核心思想是“数据切分，模型复制，梯度聚合”。其计算流程可以概括为：

  1. 数据分发：一个批次（batch）的数据被分割成 NN 个子批次（mini-batch），每个子批次被分发到一个独立的GPU上。
  2. 独立计算：
    • 在每个GPU上，进行独立的前向传播计算，产生损失。
    • 接着进行反向传播，为该GPU上分配的子批次数据计算出梯度。
  3. 梯度聚合：在完成反向传播后，所有GPU上计算出的梯度会被收集起来进行聚合操作。最常见的是 求平均，以确保所有GPU上的模型副本能基于全局批次的平均梯度进行更新。这个聚合过程通常通过 All-reduce 操作实现。
  4. 模型更新：聚合后的梯度被用于更新每个GPU上的模型副本。由于所有GPU接收到相同的平均梯度，它们的模型副本将保持同步。

图片内容描述： 原始文档中有一张名为“数据并行”的示意图，展示了数据并行的计算流程。图中包含四个并行的GPU设备（GPU 0, GPU 1, GPU 2, GPU 3），每个GPU接收一个数据子批次（B0, B1, B2, B3）。图示箭头表明数据子批次首先进入各自的GPU进行前向传播和反向传播，然后每个GPU独立计算梯度。最后，这些梯度通过箭头汇聚，表示进行聚合操作（如All-reduce），最终统一更新所有GPU上的模型。这形象地说明了每个GPU独立计算并最终聚合梯度的过程。 【若需查看原始图片详情，请参考原文中的“数据并行”示意图】

3.1.3 PS (Parameter Server) 架构

Parameter Server (PS) 架构 是一种经典的数据并行分布式机器学习框架，由Google在2014年提出。它将分布式系统的角色分为两类：

  • Server (参数服务器)：
    • 负责存储机器学习模型的所有参数。
    • 接收来自Worker的梯度。
    • 根据接收到的梯度对本地存储的参数进行更新 (update)。
    • 将最新参数分发给Worker。
  • Worker (工作节点)：
    • 从Server端获取 (pull) 当前最新的模型参数。
    • 使用本地或者远程节点的数据以及从Server获取的参数，计算关于训练参数的梯度 (compute)。
    • 将计算出的梯度发送 (push) 给Server端。

问题：in-cast的数据拥塞 PS架构的一个主要缺点是可能导致 in-cast数据拥塞。当大量Worker同时将各自计算的梯度“push”给Server时，Server端可能会因为接收带宽或处理能力的限制而成为瓶颈，导致数据传输拥塞，降低整体训练效率。这就像很多水流同时涌向一个细小的管道口，会造成堵塞。

3.1.4 Ring-all-reduce的数据并行架构

为了解决PS架构可能存在的单点瓶颈和in-cast拥塞问题，Ring-all-reduce 成为一种更高效的梯度聚合方式，尤其适用于GPU集群。

  • 原理：Ring-all-reduce 将梯度聚合和广播的过程分解为一系列点对点通信操作，数据在设备之间形成一个“环形”拓扑结构进行传输。每个设备既发送数据也接收数据，以分段的方式完成所有数据的归约和广播。
  • 优势：
    • 去中心化：没有中心服务器，避免了单点故障和瓶颈。
    • 负载均衡：每个设备在通信过程中都承担相似的工作量，使得流量传输更容易满足均衡需求。
    • 高效率：通过重叠计算和通信，以及优化带宽利用率，Ring-all-reduce 通常比PS架构在梯度聚合上更高效。
  • 实现方式：梯度聚合（pull）和传播（push）的过程本质上就是 AllReduce 的过程。Ring-all-reduce 通过将AllReduce拆解为 Reduce-Scatter 和 All-Gather 两个阶段来实现。

3.1.5 Multiple Ring-AllReduce架构

进一步优化数据并行中的梯度聚合，可以利用网络拓扑的特性。

  • 思想：由于数据并行的流量只是在所有设备上执行“加和”的操作（满足结合律和交换律），因此可以利用光交换机设计多个互质网络环。
  • 优势：通过在多个环形网络上并行传输，可以实现传输分流，显著增加传输效率，进一步降低梯度聚合的时间。这就像修建多条环形公路来分散交通流量。 【若需查看原始图片详情，请参考原文中的“Multiple Ring-AllReduce的数据并行架构”示意图，该图通常会展示多个交错的环形连接】

3.1.6 同步策略

在数据并行训练中，不同Worker之间参数更新的同步方式是影响训练效率和模型收敛行为的关键。

  • BSP (Bulk Synchronous Parallel) 同步：

    • 这是最严格的同步方式。每轮迭代（iteration），所有计算设备必须全部完成梯度计算、梯度聚合、更新模型、分发更新后的模型之后，才能够进行下一个iteration的梯度计算。
    • 优点：保证模型参数的一致性，收敛行为与单机训练一致，调试相对简单。
    • 缺点：存在木桶效应，整个训练速度受限于最慢的Worker（“短板”）。

  • ASP (Asynchronous Parallel) 完全异步：

    • 每个Worker独立异步地计算梯度，并更新参数服务器上的模型参数，然后立即获取最新参数并进行下一轮的梯度计算，无需等待其他Worker。
    • 优点：充分利用计算资源，避免木桶效应，在Worker之间传输代价较高时（如Geo-Distributed，广域网环境）尤其适用。
    • 缺点：
      • 模型参数不一致性：Worker可能基于过时的参数计算梯度，导致“梯度陈旧”问题。
      • 收敛速度和准确度影响：陈旧梯度可能导致训练震荡，收敛速度变慢，甚至无法收敛到最优解，或导致最终模型准确度下降。

  • SSP (Staleness Synchronous Parallel) 有限程度的异步：

    • SSP是BSP和ASP之间的一种折衷方案。它允许Worker使用略微过时的模型参数进行计算，但会限制这种“陈旧度（staleness）”在一定阈值之内。
    • 如果某个Worker的参数版本落后超过预设的阈值，它会被暂停，直到参数版本追上为止。
    • 优点：在提高训练速度的同时，通过限制陈旧度来缓解ASP的收敛问题，兼顾效率和准确性。
    • 缺点：需要精心选择陈旧度阈值，过大的阈值会影响收敛，过小的阈值会重新引入部分同步开销。

异步性对训练速度与准确度的影响：

  • 训练速度：异步性越高（从BSP到ASP），Worker之间的等待时间越少，理论上训练速度越快，因为没有等待其他Worker的开销。
  • 准确度/收敛性：异步性越高，模型参数的不一致性越大，训练过程中梯度可能越陈旧。这会导致：
    • 收敛速度下降：模型可能需要更多的迭代才能收敛，甚至无法收敛到相同质量的解。
    • 模型准确度下降：最终训练出的模型性能可能不如同步训练。
  • “训练越快，训练越慢！”：这句看似矛盾的话，深刻地揭示了异步训练的一个核心问题。虽然单步迭代可能更快（物理时间），但由于参数更新的混乱，模型可能需要更多的迭代才能达到相同的精度（迭代次数更多），或者无法达到同样高的精度，从有效收敛的角度看，反而“慢”了。

3.1.7 数据并行的优缺点

  • 优点：
    • 实现简单：相对于模型并行，数据并行的实现和调试通常更直接。
    • 加速训练：通过并行处理数据，可以显著减少单个Epoch的训练时间。
    • 高效率：在计算密集型任务中，数据并行能够有效利用多设备资源。
  • 缺点：
    • 模型内存没有下降：每个GPU上都需要存放模型的完整副本。这意味着如果模型本身过大，单个GPU内存无法容纳，数据并行就无法直接使用。它主要用于加速计算任务，而非解决模型内存限制问题。
    • 批次大小限制：虽然总的有效批次大小可以很大，但单个GPU上的迷你批次大小可能受限于可用内存，这会影响模型训练的稳定性或收敛性。
    • 通讯开销：梯度的聚合（All-reduce）会引入可观的通信开销，尤其是在设备数量众多或网络带宽有限的情况下。

【“3.2 模型并行 (Model Parallelism)” 和 “3.3 管道并行 (Pipeline Parallelism, PP)“的学习笔记：】

3.2 模型并行 (Model Parallelism)

模型并行 (Model Parallelism) 是一种分布式训练技术，其核心思想是将一个过大的模型在不同的计算设备（如GPU）之间进行切分，而不是像数据并行那样复制模型。它主要用于解决以下问题：

  • 核心原因 1：模型在一个计算设备上放不下。 随着深度学习模型规模的爆炸式增长（如GPT-3、Switch-C等），单个GPU的内存往往不足以容纳整个模型的参数、激活值、梯度和优化器状态。模型并行通过将模型拆分到多个设备上，使得每个设备只需存储模型的一部分，从而克服了单设备内存限制。
  • 核心原因 2：在batch_size之外进一步切分计算粒度：hiddensize上切分的可能性。 除了数据（batch_size）维度上的并行，模型并行还可以在模型内部的维度（如隐藏层大小 hiddensize 或序列长度 sequence length）上进行切分，以实现更细粒度的并行化。

主要类型： 模型并行根据切分方式可以分为两种主要类型：

  1. 层间切分 (Inter-layer Parallelism)：也被称为管道并行 (Pipeline Parallelism)。它将模型的不同层分配到不同的设备上。例如，GPU A处理模型的第1-3层，GPU B处理第4-6层，以此类推。
  2. 层内切分 (Intra-layer Parallelism)：也被称为张量并行 (Tensor Parallelism)。它在一个模型层内部进行切分，例如将一个大型矩阵乘法操作分解成多个较小的矩阵乘法，分别在不同的设备上执行。这在Transformer架构中对MLP和自注意力模块的矩阵运算尤为常见。

3.3 管道并行 (Pipeline Parallelism, PP)

3.3.1 定义与原理：层间切分

管道并行 (Pipeline Parallelism, PP) 是一种模型并行策略，通过垂直切分模型（即按层切分），将模型的不同层或一组连续的层分配到不同的计算设备（通常是不同的GPU）上。

  • 原理：每个设备只负责模型的一部分层。数据（通常是mini-batch）在这些设备之间按顺序流动，如同流水线一样。当一个设备完成其负责层的计算后，将中间结果（激活值）传递给下一个设备。
  • 目标：
    • 减少模型内存占用：这是PP最直接的好处。与数据并行每个GPU都需要完整的模型副本不同，管道并行将模型拆分到不同的GPU/主机之间，显著减少了每个设备上的模型内存占用。
    • 允许训练超大模型：使得内存无法容纳在单个设备上的巨大模型得以训练。

3.3.2 PP与DP的结合：PP+DP

管道并行（PP）和数据并行（DP）可以结合使用，形成一种混合并行策略（如GPipe等）。

  • DP的内存占用：每个GPU上都要放置完整的模型副本，因此DP在“模型层面上”并不会降低内存占用。
  • PP的内存占用：模型被拆分到不同的GPU/主机之间，从而减少了每个设备的模型内存占用。
  • PP+DP：在DP的基础上，每个DP副本（即每个DP组）内部再使用PP来切分模型。这样，既能利用DP扩展到更大批次数据，又能利用PP承载更大的模型。这意味着对于每个DP副本，其内部的各个GPU之间会交换梯度更新信息。

3.3.3 Pipeline概念与Microbatch

为了解决简单的层间切分可能导致的GPU空闲问题（当前设备计算完其部分后，需要等待下一设备完成计算才能继续反向传播），管道并行引入了微批次（Microbatch） 的概念和流水线调度机制。

  • Microbatch的定义与作用：

    • Microbatch 是比普通迷你批次（Mini-batch）更小的批处理运算块。
    • 它将一个Mini-batch进一步拆分成若干个Microbatch。
    • 作用：这些Microbatch分块进入模型管道，使得每个GPU可以同时处理不同Microbatch在不同阶段的计算。例如，当GPU 1处理Microbatch 1的前向传播时，GPU 0可能已经在处理Microbatch 2的前向传播，或者处理Microbatch 0的反向传播。这种方式提高了GPU的利用率，减少了空闲时间。

  • Batch size, Mini-batch size, Micro-batch size：

    • Batch size：机器学习算法要求的一个完整批次的大小（例如，整个模型的有效训练批次大小可能是2048）。
    • Mini-batch size：通过数据并行（DP）后，每个模型副本分到的批处理大小（例如，如果DP=2，则mini-batch_size = 2048 / 2 = 1024）。
    • Micro-batch size：继续通过管道并行（PP）后，每个GPU设备在单个计算阶段实际运行的批处理大小（例如，如果mini-batch_size = 1024，PP切分了4个阶段，micro-batch_size = 1024 / 4 = 256）。

  • Stage number与Microbatch数量：

    • 我们把模型拆分的划分数量称为 stage number（K）。
    • Microbatch的数量（M）可以不等于 stage number，通常M会远大于K，以减少“气泡”（bubble）。
    • 图片内容描述：原始文档中有一张“Pipeline in PP”示意图，展示了minibatch被拆分成microbatches，并在GPU之间以流水线方式处理的过程。图中清晰地描绘了不同GPU在不同时间点处理不同microbatch的前向传播（FP）和反向传播（BP）阶段。其中有一些“气泡（Bubble）”表示GPU的空闲时间。图示下方文字说明“层间通讯开销隐藏在计算开销之后”，并且“当网络状况极差时可能需要特别的并行调度设计”。 【若需查看原始图片详情，请参考原文中的“Pipeline in PP”示意图】

3.3.4 系统性能分析

管道并行虽然提高了GPU利用率，但仍存在其固有的性能开销。

  • Bubble时间（空闲时间比例）：

    • 定义：在管道并行中，由于设备间的依赖性（例如，一个设备必须等待前一个设备完成前向传播才能开始自己的前向传播，或者等待后一个设备完成反向传播才能开始自己的反向传播），会出现GPU空闲的时间段，这被称为“气泡”（Bubble）。
    • 影响：Bubble越大，GPU的利用率越低，训练效率越差。
    • 分析：理论上，气泡最长处近似于一个microbatch的前向传播（FP）和反向传播（BP）的时间之和。对于K个阶段和M个microbatch，Bubble time（idletime/idea time） 可以近似表示为 O(K−1M)O(\frac{K-1}{M})。
    • 优化：当K（阶段数）较小，M（microbatch数）足够大的时候，bubble time的比例会变得可以忽略不计。
    • 权衡：然而，M过大（意味着micro-batch size = N/M太小）会导致GPU计算资源无法充分利用，因为小批次数据可能无法有效地填充GPU的计算单元。好的实践是选择能够刚好让加速设备计算满载的microbatch size。

  • Peak activation memory（峰值激活内存）：

    • 激活值是前向传播的中间结果，在反向传播时需要重新计算梯度。为了节省计算，通常会将激活值存储起来。
    • 在管道并行中，峰值激活内存包括：
      • 当前正在运行的microbatch的反向传播（BP）展开部分，其内存开销约为 O(L/K⋅N/M)O(L/K \cdot N/M)，其中L是总层数，K是阶段数，N是minibatch size，M是microbatch数量。
      • 其他microbatch尚未展开的部分，其内存开销约为 O(N)O(N)。

3.3.5 PipeDream：进阶的Pipeline设计

PipeDream 是一种更进一步的管道并行设计，其核心特点是在管道并行的层面上进行异步的更新，从而最大化GPU的利用率和运行时间。 这与传统的同步管道并行（如GPipe）有显著区别。

  • PipeDream与GPipe的对比：

    • GPipe (Google Pipeline)：所有microbatch在通过整个管道完成前向和反向传播后，才进行一次梯度更新（pipeline Flush）。这种同步更新确保了模型参数的一致性，但会引入较大的气泡。
    • PipeDream：每个microbatch完成其反向传播后就马上更新模型梯度，而不是等待整个minibatch所有microbatch都完成。这意味着没有“Flush”的情况，模型是异步更新的。例如，当microbatch #5进入管道时，它可能会使用已经被microbatch #1的梯度更新过的新模型参数。

  • Per stage调度策略：

    • PipeDream每个stage计算哪个microbatch的前向传播（FP）还是反向传播（BP）的优先级由其当前活跃的microbatch决定。
    • 存在一个最优的活跃microbatch数量：理想情况下，所有worker都正在处理所有的microbatch，这样刚好不浪费memory。
    • 当管道达到稳态后，每个FP操作会为管道带来一个新的microbatch，而每个BP操作会从管道送走一个完成的microbatch。
    • 如果所有worker都在忙，处理速度为 O(workers)O(\text{workers})。
    • 最优工作microbatch数量：如果每个stage只有一个device工作，对于一个input stage来说，最优的工作microbatch数量等于 (# workers)。例如，stage 2的最优处理数量取决于它往后所有worker的处理能力。

  • 1F1B调度模式 (One-Forward-One-Backward)：

    • Pipedream在各个stage进入稳态后，会采用1F1B调度模式。这意味着在释放一个microbatch（BP完成）后，会立即处理下一个microbatch（FP开始），以保持活跃microbatch数量为最优大小。
    • 图片内容描述：原文提供了一个示例，展示了worker 3在等待3个microbatches的BP完成后，才接收下一个microbatch，以确保其运行中的microbatch数量不超过2。 【若需查看原始图片详情，请参考原文中的“Example：worker 3等到BP3个microbatches出去再接受一个，保证在运行的microbatch不超过2！”】

  • Weight Stashing：

    • 必要性：由于PipeDream每个microbatch会马上更新模型参数，而反向传播（BP）需要对应前向传播（FP）时的激活值和模型权重。为了解决这个问题，PipeDream需要“Weight Stashing”，即存储了每个microbatch在前向传播时所使用的模型参数副本。
    • 代价：这会增加内存开销。

  • Vertical Sync：

    • 由于异步更新，不同microbatch在前向和反向传播过程中可能使用了不同版本的模型参数，这被称为模型版本不一致性。
    • Vertical Sync 旨在要求一个microbatch在不同阶段（stages）使用的是一个相同时间版本的模型。
    • 实现：例如，通过n-bounded SSP（Staleness Synchronous Parallel）限制模型版本差异。
    • 代价：这可能意味着更大的存储需求或引入额外的同步开销。
    • 图片内容描述：原文中“5-FP在各个stage分别面对的是1、2、3、4BP之后的模型参数。使用vertical sync则一直使用1BP之后之后模型参数，但是意味着更大的存储需求。”展示了异步更新可能导致的不同模型版本使用情况，并说明了Vertical Sync如何通过牺牲存储来保持版本一致性。 【若需查看原始图片详情，请参考原文中的“PP上异步的更新方式”示意图】

  • PipeDream的优缺点：

    • 优势：
      • 高效并行机制：相比于数据并行，通常具有更低的机间通讯量，因为它传输的是激活值而非梯度。
      • 高GPU利用率：通过异步更新和流水线调度，最大化GPU的运行时间，减少气泡。
    • 问题/缺点：
      • 异步训练破坏训练语义：由于模型参数的异步更新，可能导致“梯度陈旧”问题，从而影响模型的收敛速度和最终准确性。这与数据并行中的ASP相似。

3.3.6 PipeDream-FLUSH

PipeDream-FLUSH 是PipeDream的一个变种，旨在解决PipeDream的两个主要问题：

  1. 问题1：Weight Stashing：PipeDream为每一个microbatch提供自己的Weight Stashing，增加了内存负担。
  2. 问题2：异步训练的准确性损失：异步更新可能导致模型收敛困难或准确度下降。
  • 核心思想：PipeDream-FLUSH 采用了1F1B调度的FLUSH版本，即在一个minibatch的所有microbatch都通过管道完成FP和BP后，才进行一次梯度更新。
  • 优点：
    • 不需要Weight Stashing：因为模型参数在整个minibatch的反向传播完成后才更新，所以所有BP阶段都使用相同的模型版本，无需存储FP时的参数副本。
    • 无需考虑model version：解决了异步更新导致的模型版本不一致问题，保证了训练语义的同步性。
    • 节约内存：中间结果保存开销逐级下降，最多不超过stage数量（第一个stage），最少为1（最后一个stage）。而原始的Pipeline Parallelism（如GPipe）每个stage的中间结果都为 num_microbatch。Less micro-batch size -> less peak memory consumption.
    • Micro-batches no longer consume memory after BP, with more micro-batches, the number of active micro-batches in the pipeline decrease.
  • 缺点：
    • 不减少Bubbletime：PipeDream-FLUSH虽然解决了异步问题，但其本质上回到了同步更新，因此并没有减少GPipe那样的气泡时间，其Bubble time开销仍然是 O(K−1M)O(\frac{K-1}{M})。
    • Bubbletime分析：看最后一个Device计算。其总时间为 M×(FP+BP)+(K−1)×FP+(K−1)×BPM \times (FP + BP) + (K-1) \times FP + (K-1) \times BP。 【若需查看原始图片详情，请参考原文中的“Bubbletime分析”公式】

3.3.7 1F1B with interleaved stages

针对管道并行中高额的Bubble time（有时甚至占比高达50%）的问题，1F1B with interleaved stages 被提出作为一种解决方案。

  • 问题回顾：

    • Pipeline FLUSH (如GPipe, PipeDream-FLUSH)：同步更新，确保准确性，但Bubble time占比非常大。
    • 不做FLUSH (如PipeDream)：异步更新，训练效率高，但会影响收敛速度和模型准确性。
    • 增加microbatch number：可以减少Bubble time，但可能导致GPU利用率下降（因为micro-batch size过小）。

  • 解决方案：1F1B with interleaved stages 的核心思想是让每个设备负责多个逻辑阶段（stages）。

    • 原理：如果逻辑上的Stage数量（K）多于物理设备数量（GPUs），那么可以将多个Stage分配给同一个物理设备。例如，模型有8个逻辑Stage，但只有4个GPU，那么每个GPU将负责2个Stage。这样，一个microbatch会在同一个物理设备上反复运行多次，处理其负责的多个逻辑Stage。
    • 图片内容描述：原文图示“Stage多于设备，每个设备负责多个stages，因此每个microbatch会反复跑各个设备多次。” 例如，FP逻辑上有8个stage，要在4个device上面循环两次。 【若需查看原始图片详情，请参考原文中的“1F1B with interleaved stages”示意图】

  • Bubble time分析：

    • 假设每个Device分配 vv 个stage。
    • 一个microbatch在一个设备上每次花费的时间为 (FP+BP)/v(FP + BP) / v。
    • 新的Bubble time 可以显著降低，约为 O(K−1vM)O(\frac{K-1}{vM})。相比于 O(K−1M)O(\frac{K-1}{M}) 减少了 vv 倍。 【若需查看原始图片详情，请参考原文中的“Bubble time分析”公式】

  • 优缺点：

    • 优点：
      • 减少 vv 倍的Bubble time：显著提高了GPU的利用率。
      • 对GPU利用率影响不大：由于拆分的是串行运算序列，只要microbatch size设置得当，可以保持较高的GPU利用率。
    • 缺点：
      • 增加 vv 倍的通讯开销：由于一个microbatch在同一个物理设备上可能需要多次进行stage间的输入输出，这会增加设备内部或设备之间的通讯频率和总量。
    • 适用场景：当网络带宽大、而计算利用率不高时，采用此方案可以有效提升整体效率。

【“3.4 张量并行 (Tensor Parallelism, TP)"、“4. 分布式大模型训练系统 - 2”、“4.1 序列并行 (Sequence Parallel, SP)"、“4.2 3D并行 (PTD-P: PP + TP + DP)“的学习笔记：】

3.4 张量并行 (Tensor Parallelism, TP)

3.4.1 定义与原理：层内切分

定义： 张量并行（Tensor Parallelism, TP）是一种层内并行（Intra-layer Parallelism）的分布式训练策略。它的核心思想是对模型中的单个层（例如，Transformer模型中的线性层或自注意力模块）的权重张量进行切分，并将这些切分后的张量分布到不同的计算设备（如GPU）上进行计算。 与数据并行（DP）复制整个模型不同，也与管道并行（PP）按层切分模型不同，TP是在一个层的内部进行细粒度的并行化。

原理： TP通过对模型权重矩阵进行分块，使得每个设备只负责计算部分矩阵乘法。由于张量并行需要在层内进行切分和计算，因此它通常需要非常频繁的设备间通信来聚合中间结果或同步梯度，这意味着它对网络连接的速度要求极高。目前主流的TP实现是专门针对Transformer架构优化的，主要应用于其核心组件：MLP（多层感知机）和自注意力模块。

3.4.2 MLP的张量并行

MLP层通常由两个线性变换（矩阵乘法）和一个非线性激活函数组成，例如 Y=G(XA1)A2Y = G(XA_1) A_2。在TP中，我们可以选择不同的方式来切分这些矩阵。

3.4.2.1 按行切分与按列切分

假设一个矩阵 AA 需要与输入 XX 相乘。

  • 按列切分 (Column Parallel)：

    • 将矩阵 AA 按列切分成 A1,A2,…,APA_1, A_2, \dots, A_P (其中 PP 是并行设备数量)，每个设备 ii 存储 AiA_i。
    • 每个设备 ii 计算 XAiXA_i。
    • 所有设备的计算结果 XA1,XA2,…,XAPXA_1, XA_2, \dots, XA_P 可以直接拼接得到完整的 XAXA。
    • 优点： 在第一个矩阵乘法 XA1XA_1 后，可以分别对 XAiXA_i 的结果应用非线性函数，而无需进行聚合通信。
    • 在Transformer的MLP中： 通常第一个线性层会采用列并行，因为其输出 XA1XA_1 往往会作为激活函数的输入，各部分的计算可以独立进行，最后再聚合。
    • 对应图示（如原始资料中所示）： “按照列分，然后组合”。

  • 按行切分 (Row Parallel)：

    • 将矩阵 AA 按行切分成 A1,A2,…,APA_1, A_2, \dots, A_P，每个设备 ii 存储 AiA_i。
    • 输入 XX 需要在所有设备上复制。
    • 每个设备 ii 计算 XAiXA_i。
    • 所有设备的计算结果 XA1,XA2,…,XAPXA_1, XA_2, \dots, XA_P 需要**进行聚合（通常是求和或Reduce操作）**才能得到完整的 XAXA。
    • 优点： 减少了每个设备存储的参数量。
    • 在Transformer的MLP中： 通常第二个线性层会采用行并行。由于其输入通常是前一个层（可能已经聚合或在每个设备上复制）的输出，并且需要聚合才能得到最终结果，所以需要进行Reduce操作。
    • 对应图示（如原始资料中所示）： “按照行分，然后加和”。

在两层MLP中： 为了减少通信开销，一个常见的策略是：

  • 第一个矩阵乘法（例如，输入到隐层的转换）采用按列切分。
  • 第二个矩阵乘法（例如，隐层到输出的转换）采用按行切分。 通过这种方式，可以在两个线性层之间巧妙地安排通信，使得只在关键点进行一次通信。
3.4.2.2 Forward与Backward的通讯过程

以一个简单的两层MLP为例 Z=NonLinear(XA)BZ = \text{NonLinear}(XA)B，其中 AA 采用列并行， BB 采用行并行。

Forward Pass (前向传播)：

  1. 输入 XX 的分发： 对于TP的第一层，输入 XX 需要拷贝到所有Worker（GPU）上，以便每个Worker能够独立计算其负责的部分。
  2. 第一个矩阵乘法 (XAXA)：
    • 假设 AA 被列切分成 A1,A2,…,APA_1, A_2, \dots, A_P。
    • 每个Worker ii 计算 Zi=XAiZ_i = XA_i。
    • 此时，每个Worker ii 得到了 XAXA 的一部分。为了应用非线性函数，这些部分需要被组合起来。然而，如果非线性函数是元素级的（如ReLU），我们可以推迟聚合，让每个Worker先对自己本地的部分 ZiZ_i 应用非线性函数得到 NonLinear(Zi)\text{NonLinear}(Z_i)。
    • 通讯（若有）： 通常这里会进行一次All-Gather来聚合所有 ZiZ_i。但如果是非线性函数，有时可以直接进行下一步的计算，再考虑聚合。原始资料中提到 “g的forward：每个worker上的forward的计算完毕，取得Z1和Z2后，GPU间做一次all-reduce，相加结果产生Z”，这里假设的 gg 是第二个线性层，而 ff 是第一个线性层。这暗示了两个线性层之间的通信模式。
  3. 第二个矩阵乘法 (ZBZB)：
    • 假设 BB 被行切分成 B1,B2,…,BPB_1, B_2, \dots, B_P。
    • 每个Worker ii 计算 ZBiZ B_i。
    • 由于 BB 是按行切分的，每个Worker ii 存储 BiB_i。为了得到完整的 ZBZB，每个Worker计算完 ZBiZB_i 后，需要进行一次All-Reduce（求和操作），将所有Worker的 ZBiZB_i 加起来，以得到完整的 ZBZB。这个聚合的结果会在所有Worker上可用。

Backward Pass (反向传播)：

  1. 梯度 dZdZ 的分发： 反向传播开始时，梯度 dZdZ（来自后续层）会作为输入。
  2. 第二个矩阵乘法（BB 的梯度计算）：
    • 计算 dBdB（相对于 BB 的梯度）和 dX′dX'（传递给前一层的梯度）。
    • 由于 BB 是按行切分的，每个Worker只负责其本地 BiB_i 的梯度计算。对于传递给前一层的梯度，需要聚合。
    • 通讯： 通常会进行一次All-Reduce来聚合所有Worker计算的梯度 dX′dX'，确保每个Worker拥有完整的梯度，以便其计算第一层 AA 的梯度。
  3. 第一个矩阵乘法（AA 的梯度计算）：
    • 计算 dAdA（相对于 AA 的梯度）。
    • 由于 AA 是按列切分的，每个Worker只负责其本地 AiA_i 的梯度计算。
    • 通讯（若有）： 对于传递给更前一层的梯度，同样可能需要All-Reduce。

总结： 当多个张量并行模块连接时，前向传播和反向传播中，通常涉及两次All-Reduce操作。例如，在一个Transformer层中，自注意力模块和MLP模块都会有各自的TP通信，通常整体包含Forward的两次All-Reduce和Backward的两次All-Reduce。

3.4.3 Self-Attention的张量并行

自注意力机制也包含多个线性变换（用于计算Q、K、V），因此同样可以应用张量并行。

  • 基于不同注意力头 (K_i, Q_i, V_i) 的切分：
    • 在自注意力机制中，通常有多个注意力头（Multi-Head Attention）。TP可以沿着注意力头的维度或隐藏维度进行切分。
    • 例如，可以将Q、K、V矩阵按列切分，使得每个设备只负责计算部分注意力头的Q、K、V。
    • GEMMs (General Matrix Multiplications)： 这些切分操作最终都会归结为通用矩阵乘法，可以高效地并行执行。
  • 输出Y的连续性：
    • 自注意力模块的输出 YY 通常是完整的子矩阵部分，这意味着它可以直接作为下一个模块（例如MLP的第一个线性层）的输入，并继续进行按行切分的张量并行，无需额外的通信。
  • 通讯： 与MLP类似，自注意力模块的张量并行通常也需要总共两次All-Reduce操作（一次在前向，一次在反向）来聚合中间结果。

3.4.4 张量并行的优缺点

优点：

  • 减少每个GPU所需的内存量： 由于模型权重被切分，每个GPU只需存储部分权重，从而显著降低了单个GPU的模型内存占用，使得训练超大模型成为可能。
  • GPU利用率保持较高 (无Bubble)： 与管道并行不同，张量并行通常不会引入“气泡时间”（bubble time），因为它是在层内部进行并行计算，且通信通常可以与计算重叠，因此GPU的利用率可以保持较高水平。

缺点：

  • 非常频繁的同步 (AllReduce)： 张量并行需要在层内进行多次All-Reduce操作来聚合中间结果或梯度。这意味着它对网络带宽和延迟要求极高。
  • 高吞吐量的网络连接： 为了保持高吞吐量，需要极快的网络连接，例如NVLink或InfiniBand等高速互联技术，通常这限制了TP只能在单个节点内部或高速集群的“超节点”之间使用。如果网络带宽较低，频繁的通信会导致明显的性能瓶颈。

4. 分布式大模型训练系统 - 2

这个部分将深入探讨更高级和复合的分布式训练策略，以应对大模型训练中内存和计算效率的进一步挑战。

4.1 序列并行 (Sequence Parallel, SP)

4.1.1 激活内存开销问题

背景：

  • 管道并行 (PP) 的局限性： 即使使用了PP，每个Stage（尤其是靠前的Stage）仍然需要存储大量的前向传播 (FP) 临时结果（即激活值 Activation），以便在反向传播 (BP) 时进行梯度计算。这些激活值是内存消耗的大户。
  • PP切分过多： 虽然切分更多的PP Stage可以减少每个Stage的模型内存占用，但相对地，激活值的内存开销占比会变得更高，这成为新的内存瓶颈。

挑战： 如何进一步减少激活值的内存占用，尤其是在Transformer模型中？

4.1.2 TP与SP结合的提出

原始的张量并行 (TP) 方案通常只关注矩阵乘法本身的切分，而忽略了Transformer模型中其他重要的操作，如层归一化 (Layer Normalization, LayerNorm) 和 Dropout。

  • 原始TP的局限性： 在原始的TP中，这些操作通常并没有被拆分，而是在所有TP设备中重复进行计算和存储。虽然这些操作的计算量不大，但它们需要消耗大量的激活内存。

解决方案： 提出将层归一化和Dropout的计算和存储工作也在TP的设备之间进行分割，这就是序列并行 (Sequence Parallelism, SP) 的核心思想。

4.1.3 SP的原理：层归一化与Dropout的切分

核心观察 (Observation)：

  • 层归一化和Dropout的操作是独立的针对每个Token/样本进行的。这意味着对一个Token进行LayerNorm/Dropout不会影响其他Token的计算。
  • 它们需要Token/样本的完整向量才能完成操作，而不是像矩阵乘法那样对特征维度进行分解。

核心思想 (Idea)： 既然LayerNorm/Dropout是针对每个Token独立的，我们可以利用这一点，在TP的机器之间引入SP，让每个TP设备只处理一部分Token的LayerNorm/Dropout操作。

具体做法：拆解All-reduce操作 为了实现这种Token维度的并行，SP将TP中通常用于聚合的All-Reduce操作进行拆解和重新设计：

  • Forward Pass (前向传播)：
    • 在TP计算之前，进行 Token-level Reduce-scatter (或类似操作) 将不同Token的数据分发到不同的TP设备上。
    • 在TP计算之后，进行 All-gather 将所有TP设备的输出重新聚合，以便后续操作或输出。
  • Backward Pass (反向传播)：
    • 在梯度计算之前，进行 All-gather。
    • 在梯度计算之后，进行 Reduce-scatter 将梯度分发给负责不同Token的设备。

更具体的转换：

  • 原TP：All-Reduce in FP, no operation in BP. (这里可能是指TP的一种变体，或者强调其核心通信点)
  • SP的引入：
    • all-gather in the FP, and reduce-scatter in the BP. (一种实现方式，用于聚合输入或分发梯度)
    • reduce-scatter in FP, and all-gather in BP. (另一种实现方式，用于分发输入或聚合梯度)

理解 g 聚合 token 的作用：

  • 原始资料中提到 g 负责收集所有的token，然后执行TP的子矩阵运算 和 g 负责在做reduce的时候，让每个TP的机器只拿去一部分tokens的reduce结果，然后进行各自独立的中间操作。这表明在TP+SP的组合中，可能存在一个“门控”或聚合/分发机制 g，它在FP阶段聚合所有Token以便进行TP的矩阵运算（因为矩阵运算需要整个特征维度），而在BP阶段只将一部分Token的Reduce结果分发给TP设备，以便各自独立处理。

4.1.4 TP + SP在MLP中的应用

  • SP在Token维度上聚合/分割： 每个TP设备负责处理输入序列中一部分Token。
  • TP在特征维度上聚合/分割： 每个TP设备负责处理模型权重矩阵中一部分特征。

结合效果： 通过这种结合，可以同时在Token维度和特征维度上进行并行，进一步降低单个设备的内存负担。

4.1.5 Activations memory分析

通过SP，可以将LayerNorm/Dropout等操作的激活值也分布到不同的TP设备上，从而显著减少每个设备存储的激活值总量。这意味着：

  • TP解决模型权重本身的内存问题。
  • SP解决TP之外的、在每个设备上重复存储的激活值内存问题。

总结： SP是对TP的一种补充，它通过并行化LayerNorm和Dropout等操作来减少Transformer模型中的激活内存开销，使得在有限内存的设备上训练更大的模型成为可能。

4.2 3D并行 (PTD-P: PP + TP + DP)

4.2.1 引入更多并行策略的原因

背景： 随着大模型规模的指数级增长，单一的并行策略往往会遇到瓶颈，无法满足训练需求。

  • 只使用DP (数据并行)：
    • 优点： 擅长分割计算任务，理论上可以无限扩展，对模型本身内存占用无影响。
    • 缺点： 不擅长分割模型以减少内存开销。每个GPU都需要存储完整的模型副本，当模型过大时，单卡内存会成为瓶颈。
  • 只使用TP (张量并行)：
    • 优点： 擅长分割模型/计算，可以显著降低单个GPU的模型内存占用。
    • 缺点： 需要非常大量的设备间通信开销（频繁的All-Reduce），在低带宽或高延迟的网络环境下会导致明显的性能瓶颈。通常限制在高速互联的节点内部。
  • 只使用PP (管道并行)：
    • 优点： 擅长分割模型/计算，可以显著降低单个GPU的模型内存占用。
    • 缺点： 存在“气泡时间” (bubble time)，导致GPU利用率降低，并行效率稍弱。

结论： 单一并行方式的扩展能力存在上限，受限于不同的维度（如 batch_size, hidden_size, layer_size）。

3D并行的提出： 将数据并行 (DP)、张量并行 (TP) 和 管道并行 (PP) 结合起来，可以独立扩展每个维度的并行度，从而在训练上百或上千个GPU的超大模型时实现更灵活、更强的可扩展性。

4.2.2 3D并行组合示例

3D并行可以根据具体的模型、集群资源和训练需求，灵活地组合DP、TP和PP的并行度。

  • DP=2, TP=2, PP=2： 假设有8个GPU，可以分成2个DP组，每个DP组有4个GPU，这4个GPU又分成2个TP组，每个TP组的2个GPU再进行2层的PP。
  • DP=4, TP=4, PP=4： 这代表了极高的并行度，通常需要大规模的GPU集群。
  • DP=4, TP=2, PP=4 或 DP=2, TP=4, PP=4： 这些都是不同的组合方式，用于在特定场景下优化性能。

关键问题： 如何确定在给定GPU数量下，各个并行策略的程度是多少？这通常需要复杂的搜索和优化。

4.2.3 通讯开销分析

在3D并行中，通信开销是性能的关键影响因素，需要对PP和TP的通信量进行量化分析。 假设：

  • MM: microbatch size（微批次大小）
  • SS: sequence length（序列长度）
  • hh: embedding/activation size（嵌入/激活维度）
  • PTPP_{TP}: TP并行度（TP replicas数量）
  • PPPP_{PP}: PP并行度（PP stages数量）

  1. 管道并行 (PP) 通讯：

    • 对于每对连续的设备（在PP流水线中相邻的Stage），在前向传播和反向传播中，每次Microbatch传输的通信量大约为 M⋅S⋅hM \cdot S \cdot h。
    • 这个通信量是传输激活值（中间结果）。
    • 计算： 每对连续设备每次微批次的前向+反向通信量为 2×(M⋅S⋅h)2 \times (M \cdot S \cdot h)。
    • 特点： PP的通信量主要取决于输入/输出的激活值大小，相对较少。

  2. 张量并行 (TP) 通讯：

    • 在TP中，涉及模型参数或梯度的All-Reduce操作。
    • 对于每个Transformer层，需要对大小为 M⋅S⋅hM \cdot S \cdot h 的张量在 PTPP_{TP} 个模型副本（TP组内的设备）之间进行两次All-Reduce操作（一次在前向，一次在反向）。
    • 每次All-Reduce的通信量大约是 2×(M⋅S⋅h)×PTP−1PTP2 \times (M \cdot S \cdot h) \times \frac{P_{TP}-1}{P_{TP}} (这通常被简化为 2×(M⋅S⋅h)2 \times (M \cdot S \cdot h)，因为All-Reduce的理论最优通信量是 2×data_size×N−1N2 \times \text{data\_size} \times \frac{N-1}{N} 或 2×data_size2 \times \text{data\_size} ）。
    • 在一个层中，前向和反向各两次All-Reduce，总共是 4×2×(M⋅S⋅h)4 \times 2 \times (M \cdot S \cdot h)。
    • 计算： 每个Microbatch在每个TP组内的每个设备，每个层，总通信量约为 8⋅M⋅S⋅h⋅PTP−1PTP8 \cdot M \cdot S \cdot h \cdot \frac{P_{TP}-1}{P_{TP}}。
    • 影响： 如果一个设备负责 kk 层模型（例如在PP中一个Stage负责多层），那么TP的通信量将增加 kk 倍。而PP的通信量则不受 kk 的影响（因为只发生在Stage之间）。

总结：

  • TP的通信量通常远大于PP，且对 PTPP_{TP} 敏感。
  • PP的通信量相对较小，但会引入气泡时间。

4.2.4 各并行策略在3D并行中的作用

  • TP + PP 的核心作用： 它们主要用于切分模型本身，以使模型能够被装载到有限内存的设备上，并允许训练得以正常进行。
  • DP 的核心作用： 在TP+PP已经切分好模型的基础上，DP可以不断增加模型副本（在DP组之间），从而扩展到更大规模的分布式训练上，提高整体吞吐量。DP并不直接降低单卡模型内存，但其通信效率相对较高。

GPU利用率与并行细粒度：

  • GPU本身是大规模并行计算设备。
  • 如果分布式切分做得过于细致，导致 microbatch size 过低，即使可以降低气泡时间（如PP优化），也可能导致GPU利用率不足。这是因为过小的 microbatch 无法充分填充GPU的计算单元，导致计算效率下降。
  • 因此，在实际应用中，需要在内存节省、通信开销和GPU利用率之间找到一个最佳平衡点。

4.3 全切片数据并行 (Fully Sharded Data Parallel, FSDP / ZeRO)

4.3.1 引入背景：现有并行策略的局限性

在处理大规模模型训练时，传统的并行策略会遇到各自的局限性：

  • 管道并行 (Pipeline Parallelism, PP)：虽然能有效切分模型以减少单个设备的内存占用，但其存在“bubble”问题（即部分GPU空闲时间），降低了计算效率。
  • 张量并行 (Tensor Parallelism, TP)：擅长在层内切分模型，但需要非常大量的设备间通讯开销。在低带宽（如跨节点）的网络条件下，这会导致明显的传输延迟，因此通常仅适用于高速连接的计算设备之间（如单机内的NVLink）。
  • 数据并行 (Data Parallelism, DP)：通讯量相对适中，并且可以通过按层传输等技术隐藏大量传输开销，不存在PP的“bubble”问题。然而，DP的缺点在于每个设备都需要存放完整的模型副本，这导致存储资源大量重复，无法有效减少单个模型副本所需的内存。当模型参数巨大时，即使是DP也可能因模型本身无法放入单个设备内存而失效。

因此，需要一种新的并行策略，它既能像DP一样高效利用计算资源，又能像模型并行一样分担模型的内存占用。

4.3.2 ZeRO (Zero Redundancy Optimizer) 核心思想

ZeRO（零冗余优化器）是一种旨在消除数据并行中内存冗余的优化技术，其核心思想是让DP的各个副本之间也能均摊模型状态的内存负担。它通过**“按需收集，迅速分发 (On-demand gather, immediate scatter)”的策略实现。 具体而言，每个设备只保存一部分参数、梯度或优化器状态**。当需要进行计算（如前向或反向传播）时，设备会按需从其他设备收集所需的完整部分；一旦计算完成，便立即丢弃这些临时收集到的数据，只保留自己负责的“切片（slice）”，并将计算结果（如梯度）分发给负责保存对应参数的设备。

4.3.3 训练内存开销分析：Model states与Residual states

在大模型训练中，所需的内存远大于模型参数本身的大小。内存开销主要分为两大部分：

  • Model states（模型状态）：
    • 模型参数 (Model Parameters)：即模型权重。
    • 模型梯度 (Model Gradients)：反向传播计算出的参数更新方向。
    • 优化器参数 (Optimizer Parameters)：例如，Adam优化器需要存储动量（momentum）和方差（variance）两个状态，通常各占用与模型参数相同大小的内存。
  • Residual states（残余状态）：
    • 中间激活值 (Intermediate Activations)：前向传播过程中每一层输出的激活值，在反向传播时需要重新计算或存储。
    • 临时缓存 (Temporary Buffers)：各种临时计算和存储所需空间。
    • 未使用的内存片段 (Unused Memory Fragments)：内存碎片等。

混合精度训练下的内存占用：

  • 假设模型参数有NN个值。
  • 模型参数和梯度通常以FP16（半精度浮点数）保存，各占2N2N字节。
  • 优化器参数是内存占用的大头：Adam优化器需要存储FP32（单精度浮点数）的模型参数副本（4N4N字节），以及FP32的动量（momentum, 4N4N字节）和方差（variance, 4N4N字节）。
  • 总计：2N2N (FP16参数) + 2N2N (FP16梯度) + 4N4N (FP32参数副本) + 4N4N (动量) + 4N4N (方差) = 16N16N 字节（原始资料中提到的是12N12N字节，可能是忽略了FP16参数本身，但核心在于优化器状态占据了绝大部分）。

例如，一个15亿参数的GPT-2模型，即使参数本身仅需要3GB（FP16），也无法在32GB内存的单个GPU上训练，因为优化器状态等会迅速耗尽内存。

4.3.4 ZeRO的Partition类型与阶段

ZeRO根据模型状态切分的程度，分为三个阶段（ZeRO-1, ZeRO-2, ZeRO-3）：

  • ZeRO-1: Optimizer State Partitioning
    • 切分内容：只切分优化器状态（如Adam的动量和方差）。
    • 通信模式：
      • 前向：无额外通信。
      • 反向：通过Reduce-Scatter操作分发梯度到负责对应优化器状态的设备。
      • 参数更新后：通过All-Gather将更新后的参数分发给所有设备。
  • ZeRO-2: Optimizer State Partitioning + Gradient Partitioning
    • 切分内容：切分优化器状态和梯度。
    • 通信模式：
      • 前向：无额外通信。
      • 反向：按需逐层Reduce-Scatter分发梯度给负责的设备。
      • 参数更新后：通过All-Gather将更新后的参数分发给所有设备。
  • ZeRO-3: Optimizer State Partitioning + Gradient Partitioning + Parameter Partitioning
    • 切分内容：切分优化器状态、梯度和模型参数。这是最激进的切分策略，也是FSDP的基础。
    • 通信模式：
      • 前向：按需All-Gather收集参数到当前GPU进行计算，计算完成后即丢弃非本设备负责的参数。
      • 反向：按需逐层All-Gather收集参数，按需逐层Reduce-Scatter分发梯度给负责的设备。
      • 参数更新后：无额外通信，因为每个设备只更新自己负责的参数切片，无需再分发。

4.3.5 ZeRO-3详细示例

假设一个线性层权重矩阵 WW 是 1024×40961024 \times 4096，在4个GPU上进行ZeRO-3分配。

  • 每个GPU只保存 WW 的一个切片 (slice)：

    • GPU0: W[:,0:1024]W[:, 0:1024]
    • GPU1: W[:,1024:2048]W[:, 1024:2048]
    • GPU2: W[:,2048:3072]W[:, 2048:3072]
    • GPU3: W[:,3072:4096]W[:, 3072:4096]

  • 前向计算时：

    1. 每个GPU先进行All-Gather操作，临时从其他GPU收集它们负责的切片，从而在本地得到完整的 WW 矩阵 (1024×40961024 \times 4096)。
    2. 使用完整的 WW 进行计算。
    3. 完成计算后，丢弃临时组合的完整 WW，只保留自己负责的 WW 切片。

  • 反向计算时：

    1. 每个GPU再次进行All-Gather操作，临时在本地得到完整的 WW 矩阵。
    2. 完成反向传播计算。
    3. 每个GPU只保留自己负责的梯度切片。对于非自己负责的梯度切片，通过Reduce-Scatter操作发送给对应的GPU。

  • 梯度更新：

    1. 每个GPU只更新自己负责的参数切片。

4.3.6 内存下降效果

ZeRO-3带来的内存下降是惊人的。通过将模型参数、梯度和优化器状态全部进行切分，每个GPU的内存占用显著降低，从而使得训练超大模型成为可能。

4.3.7 ZeRO的额外通讯开销分析

自然地，将模型状态切分到不同副本之间必然会导致更多的通信。然而，ZeRO的设计目标是尽量将这些通信开销隐藏或最小化。

  • 数据并行的通讯开销回顾：

    • 传统数据并行中，在每个minibatch的梯度计算完成后，会使用AllReduce操作将不同GPU上的梯度聚合并分发。
    • AllReduce通常通过Reduce-Scatter + All-Gather实现（如Ring AllReduce）。
    • 对于模型大小为 NN 的参数，其通信开销通常为  2N/B~2N/B（BB 为带宽）。

  • ZeRO-1与ZeRO-2的通讯开销：

    • 对于每一个层的梯度，ZeRO-1和ZeRO-2会使用Scatter-Reduce从所有模型副本处聚合各切片的梯度到对应的本地副本（开销  N/B~N/B）。
    • 随后，在梯度收集、更新模型并分发时，会使用All-Gather将更新后的参数收集到每一个模型副本上（开销  N/B~N/B）。
    • 可见，ZeRO-1和ZeRO-2的传输开销本质上是原始AllReduce的分阶段完成，与原始DP的传输开销近似。ZeRO-2进一步将这些操作均摊到整个训练的各层反向传播中，逐层按需完成。

  • ZeRO-3的通讯开销：

    • 由于模型参数也被切分，前向和反向传播过程都需要从其他副本All-Gather拉取模型参数。每次All-Gather一个完整模型参数的开销为  N/B~N/B。由于“按需收集，迅速分发”原则，FP完成后即丢弃非本设备参数，BP时需要再次All-Gather。因此，All-Gather整个模型两次，开销  2N/B~2N/B。
    • 反向传播结束时，需要Scatter-Reduce整个模型梯度一次，开销  N/B~N/B。
    • 总开销约为 3N/B3N/B。
    • 因此，ZeRO-3的总通信时间开销约为原始DP的1.5倍。这是一个可以接受的开销，尤其是在内存限制是主要瓶颈时。

4.3.8 Superlinearity

通过ZeRO，每个GPU释放出大量的内存。这些内存可以被用来增加Batch Size。在很多情况下，Batch Size越大，计算利用率越高，这可以带来**Superlinearity（超线性）**的性能提升，即系统性能的提升超过了硬件资源的线性叠加。

4.3.9 ZeRO与模型并行的关系对比

ZeRO和模型并行（MP）虽然都涉及将模型状态分布式存储在不同GPU上，但它们的核心侧重点和数据流动方式有所不同：

  • 模型并行：通常是“模型不动，数据动”。即模型被固定地切分到不同设备上，数据（或激活值）在设备之间流动。
  • ZeRO：更像是“数据不动，模型动”。即数据并行中，每个设备处理独立的数据子集，但为了减少冗余，模型的各个部分在需要时才被“召唤”到特定设备，使用完后即“送走”。

更高效的部分：

  • 模型可以提前移动 (Pipeline)：在计算N−1N-1层的输出时，可以提前传输NN层的模型参数，为NN层的计算做准备。
  • PP必须等待数据：管道并行（PP）必须等到N−1N-1层的计算完成，将计算结果（激活值）发送给NN层的设备后，才能进行后续计算。
  • 传输量对比：PP的传输量通常是层间中间结果（激活值），而ZeRO的传输量是模型大小。激活值的大小可能小于或大于模型参数大小，具体优劣需结合模型、训练算法、GPU性能和网络拓扑具体考量。在实践中需要进行性能测试来选择。

与张量并行的结合：

  • 张量并行（TP）和ZeRO切分的是不同维度：TP通常切分隐藏层维度（hidden_size），而ZeRO主要切分Batch Size维度上的模型冗余。
  • 在有限的Batch Size（如2048）下，两者结合可以获得更大的并行能力，从而扩展到成千上万个GPU上进行训练（例如，TP与FSDP结合）。

3D并行 + ZeRO：

  • 在TP组内可以继续使用FSDP（ZeRO-3）进行数据并行，即每个GPU只保存张量的一个切片。
  • 前向/反向传播在TP组内部使用3D并行的传输（TP、PP），同时在DP副本之间使用FSDP的All-Gather / Scatter-Reduce来收集和分发参数和梯度。这种组合能够最大化并行度，应对极端规模的模型训练。
  • 【若需查看原始图片详情，请参考原文中的“3D并行+Zero”示意图】

4.4 混合专家模型 (Mixture of Experts, MOE)

4.4.1 定义与核心思想：稀疏模型架构

混合专家模型（MoE）是一种稀疏模型架构，旨在进一步增加大模型的规模，同时控制计算成本。其核心思想是：

  • 将一个大模型拆分成多个独立的小型模型，称为专家 (expert)。
  • 在每次运行一个输入（如一个Token或一个Batch）时，只有部分的模型参数参与工作。
  • 通过这种方式，在维持模型巨大容量的同时，实际激活的计算量保持在可控范围，从而达到节省计算资源的效果。

4.4.2 SwitchTransformer中的MoE应用 (FFN)

在许多MoE模型中，例如Google的SwitchTransformer，MoE层通常被应用于前馈网络 (Feed-Forward Network, FFN) 部分。

  • Transformer模型中的FFN层是计算密集型部分，将其替换为MoE层可以显著增加模型容量，同时通过稀疏激活控制计算量。
  • 【若需查看原始图片详情，请参考原文中的“SwitchTransformer: MOE in FFN”示意图】

4.4.3 门的计算与专家选择

MoE模型引入了一个可训练的门 (gate) 机制，以确保稀疏性并根据输入动态选择专家：

  • 门的计算：对于给定的输入 xx，门会类似于分类过程，为不同的专家分配不同的概率或分数。这通常通过一个小型的前馈网络加上Softmax函数实现。
  • 专家选择：门会根据这些分数，选择前 kk 个得分最高的专家来处理当前的输入。通常 kk 值很小（例如 k=1k=1 或 k=2k=2）。
  • 结果组合：选定的专家独立地对输入进行处理。最终的输出是这些选定专家的输出的加权和，权重通常就是门计算出的专家分数。

4.4.4 专家并行中的All-to-all通讯

当MoE模型进行分布式训练时，通常会采用专家并行 (Expert Parallelism) 策略。这意味着不同的专家被放置在不同的设备上。在这种设置下：

  • 通信需求：每个输入需要被路由到其被选中的专家所在的设备。而每个设备可能需要处理来自不同输入的Tokens，这些Tokens被路由到其本地的专家。
  • All-to-all通信：这种路由过程通常通过All-to-all集合通信操作来实现。
    • 在前向传播 (FP) 时，通常有两次All-to-all通信：一次将Tokens路由到它们选定的专家，另一次将专家处理后的结果路由回原始Tokens的设备。
    • 在反向传播 (BP) 时，也有两次All-to-all通信，处理梯度流的路由。
  • 总共四次All-to-all通信：每个MoE层在FP和BP时各发生两次All-to-all通信，用于传输Tokens及其梯度。

4.4.5 Load Balancing Loss

在MoE模型中，一个常见的问题是负载不均衡：某些专家可能被过度选择，而另一些专家则很少被使用（“starving experts”）。这导致资源浪费和训练效率下降。为了解决这个问题，通常会在损失函数中增加一个负载均衡损失 (Load Balancing Loss) 项。

  • 原理：该损失项旨在鼓励门机制将Tokens均匀地分配给所有专家，防止少数专家过度活跃。
  • 计算方式：
    • Load Balancing Loss=∑i=1NFreq(i)⋅Prob(i)\text{Load Balancing Loss} = \sum_{i=1}^{N} \text{Freq}(i) \cdot \text{Prob}(i)
    • 其中，Freq(i)\text{Freq}(i) 是专家 ii 被选中的频率（或“出头鸟”的次数，代表其被分配到的Tokens数量），Prob(i)\text{Prob}(i) 是门函数给专家 ii 的平均概率值。
    • 通过最小化这个损失，模型会倾向于让所有专家被选中的概率相似，并且每个专家处理的Tokens数量也相对均匀。
  • 效果：降低损失的结果是让所有专家都“不做出头鸟”，最终达到负载均衡。这有助于提高所有专家的利用率，并稳定训练过程。

稀疏程度与训练时间：MoE模型的目标是实现高稀疏度，即在每次计算时激活最少的专家。虽然计算量减少了，但由于复杂的通信（All-to-all）和负载均衡问题，MoE模型的训练时间仍然可能面临挑战。

4.5 并行的自动化搜索

4.5.1 挑战：搜索空间与硬件考量

在实际的大模型训练中，面对多种并行策略（数据并行DP、管道并行PP、张量并行TP、序列并行SP、混合专家模型MoE等）及其组合，如何选择最优的并行策略是一个巨大的挑战。

  • 组合爆炸的搜索空间：即使是相对简单的3D并行（DP+TP+PP），当面对多达24个GPU时，其并行度（例如，DP×TP×PP=24DP \times TP \times PP = 24）就有多种组合，如 2×4×32 \times 4 \times 3、 2×2×62 \times 2 \times 6、 1×1×241 \times 1 \times 24 或 2×1×122 \times 1 \times 12 等。每种组合又可能对应不同的放置策略，使得搜索空间呈组合爆炸式增长。
  • 硬件异构性与网络拓扑：最优的并行策略不仅取决于模型结构，还与具体的硬件环境（如GPU型号、内存大小）以及网络拓扑结构（如NVLink直连、InfiniBand交换机、以太网）紧密相关。例如，张量并行需要极快的网络连接，通常适用于单机内GPU间的NVLink连接，而跨机通信则更倾向于数据并行或管道并行。因此，在选择并行策略时，需要同时考虑网络拓扑和计算设备的能力。
  • 试错成本高昂：在数百或数千个GPU的规模上，手动尝试不同的并行化策略是极其缓慢和昂贵的。每次实验都需要耗费大量计算资源和时间。

鉴于这些挑战，研究人员提出了一个核心问题：我们能否自动执行给定模型和集群的并行化过程？ 这促使了自动化并行策略搜索技术的发展。

4.5.2 Alpa：自动化Inter-和Intra-Operator并行

Alpa是一个旨在自动化Inter-Operator（操作符间）并行和Intra-Operator（操作符内）并行以实现分布式深度学习的系统。

  • Inter-Operator Parallelism (操作符间并行)：主要指管道并行 (PP)，即将模型的不同层（或一组操作符）放置在不同的设备上，形成一个流水线。
  • Intra-Operator Parallelism (操作符内并行)：主要指张量并行 (TP) 和数据并行 (DP)，即将单个操作符（如矩阵乘法）或其内部的数据进行切分，分布到多个设备上并行计算。

Alpa的工作原理和特点：

  1. 多样化的策略选择和并行方式：Alpa能够综合考虑多种并行策略，并根据模型的特性和硬件资源进行选择和组合。
  2. 网络拓扑决定Device Placement：Alpa会分析网络拓扑结构，以智能地决定设备（GPU）的放置策略，例如哪些GPU应该组成一个张量并行组，哪些应该用于数据并行，从而最小化通信开销。
  3. API层次搜索：Alpa通过在API层进行抽象和搜索，避免了直接在低级代码层面操作，从而大大简化了并行策略的表示和搜索过程。
  4. 规划方法进行搜索：Alpa采用了一种分阶段的规划方法 (planning method) 来解决复杂的并行搜索问题：
    • Inter-op Pass (操作符间通行)：这个阶段主要搜索管道并行 (PP) 策略。它关注如何将模型层间进行切分并放置到不同的设备或设备组上。通过设备放置（Device Placement）的决策，确定PP的拓扑结构。
      • 【若需查看原始图片详情，请参考原文中的“Inter-op pass: Device placement”示意图】：该图展示了Inter-op pass如何将模型（如一个Transformer block）划分为多个阶段，并分配给不同的设备或设备组。
    • Intra-op Pass (操作符内通行)：在确定了Inter-op并行策略后，这个阶段在每个设备组内部搜索张量并行 (TP) 和数据并行 (DP) 策略。它关注如何将单个操作符（如矩阵乘法）或Batch数据在设备组内部进行切分。
      • Cost Function (成本函数)：在Intra-op pass中，Alpa会使用一个成本函数来评估不同TP/DP策略的性能（例如，通信开销、计算时间、内存占用）。通过最小化这个成本函数，Alpa能找到给定设备组内的最优Intra-op并行策略。

示例：Wide-ResNet Partition on 16 GPUs Alpa通过其自动化并行化能力，能够有效地为复杂的模型（如Wide-ResNet）在特定数量的GPU（如16个）上找到高效的并行策略。

  • 【若需查看原始图片详情，请参考原文中的“Example: Wide-ResNet Partition on 16 GPUs”和“Automatic Parallelization with Alpa”示意图】：这些图示展示了Alpa如何为Wide-ResNet模型在16个GPU上进行自动化的并行切分和设备放置，包括结合Inter-op和Intra-op并行，以达到最佳性能。这种可视化有助于理解自动化并行策略是如何在实际模型中应用的。

通过这种分层和自动化的方式，Alpa极大地降低了分布式训练中并行策略选择的复杂性，使得开发者能够更高效地训练大规模深度学习模型。

MLSYS编程框架MLSYS推理服务