USP: A Unified Sequence Parallelism Approach for Long Context Generative AI

USP：一种用于长上下文生成式AI的统一序列并行方法
作者/机构: Jiarui Fang, Shangchun Zhao (Tencent)

A1 主要贡献

本文旨在解决生成式AI模型中日益增长的长上下文处理需求所带来的挑战。随着Claude、GPT-4、Gemini 1.5 Pro和Sora等模型将上下文长度推向新的极限，序列并行（Sequence Parallelism, SP）成为一项关键技术。然而，现有的主流SP方法，即DeepSpeed-Ulysses和Ring-Attention，各自存在局限性。DeepSpeed-Ulysses的并行度受限于注意力头的数量，且不适用于分组查询注意力（GQA）和多查询注意力（MQA）；而Ring-Attention因将矩阵乘法细分而导致计算效率下降。这些问题阻碍了SP技术的广泛应用。

本文的核心研究目标是提出一种更通用、更高效的SP方法，并系统性地分析其如何与现有的并行策略（数据并行、张量并行、ZeRO、流水线并行）结合，以构建一个强大的4D混合并行系统。

主要贡献包括：

  • 提出一种统一的序列并行方法（Unified Sequence Parallelism, USP）：该方法整合了DeepSpeed-Ulysses和Ring-Attention，旨在克服两者各自的缺点。USP通过将SP进程组划分为一个二维网格，分别在行和列上运行Ulysses和Ring并行，从而能够灵活适应不同的模型架构（如GQA/MQA）和硬件网络拓扑（如同构或异构网络），表现出更强的鲁棒性。
  • 系统性分析SP在4D并行中的应用并提供最佳实践：本文深入探讨了SP与张量并行（TP）、ZeRO和流水线并行（PP）结合使用时的复杂关系和性能权衡。通过对通信开销和内存成本的详细分析，本文提出了一系列最佳实践建议（Tips），指导研究人员和工程师如何在分布式Transformer计算中有效地应用SP，以实现长序列训练。

A3 背景知识：序列并行方法

Transformer模块计算基础

在深入探讨SP概念之前，我们先回顾标准Transformer模块的计算过程。本文使用的符号如表1所示，其中隐藏维度$d = h_c \times h_s$。

表1：Transformer参数符号表

给定输入序列$Q, K, V \in R^{L \times d}$，其中L是序列长度，d是隐藏维度，我们按如下方式计算输出矩阵：

每个自注意力子层之后都跟一个前馈网络（FFN），该网络以相同的方式独立地应用于每个位置。

早期序列并行探索

与张量并行（TP）【【4】Megatron-lm: Training multi-billion parameter language models using model parallelism，2019】和ZeRO【【5】Zero-offload: Democratizing billion-scale model training，2021】相比，针对Transformer模型的SP研究在很长一段时间内相对不成熟。挑战在于注意力计算的特性，其中序列维度在softmax之后是矩阵乘法中的一个公共维度，这使得在切分序列维度后很难将张量和计算分布到多个节点上。

Megatron-LM序列并行

早期的SP尝试【【6】Sequence parallelism: Making 4d parallelism possible，2021；【7】Colossal-ai: A unified deep learning system for large-scale parallel training，2023；【8】Lightseq: Sequence level parallelism for distributed training of long context transformers，2023】并不成功，常导致冗余的内存消耗【6】和低效的通信模式【8】。对于长输入序列，最佳实践是采用Megatron-LM的序列并行。该方法优化了TP的AllReduce操作，在保持相同通信开销的同时降低了激活值的内存成本。如图1所示，Megatron-LM序列并行的原理类似于ZeRO-2【5】。它将对复制张量（左图）的AllReduce操作替换为对分区数据（右图）的等效all-gather和reduce-scatter操作。由于AllReduce操作本质上是All-gather和Reduce-scatter的组合，因此通信成本保持不变。输入和输出张量的大小在N个计算设备上减少了1/N。由于输入/输出张量中的序列维度被分区，因此它被命名为序列并行。然而，这种形式的序列并行不能独立于张量并行使用，并且无论并行度如何，通信量都保持不变。


图1：Megatron-LM序列并行的原理。

现代序列并行方法：SP-Ulysses与SP-Ring

独立的SP技术的成熟以2023年末两篇里程碑式论文的发表为标志。DeepSpeed-Ulysses【【1】Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models，2023】（称为SP-Ulysses）和Ring-Attention【【2】Ring attention with blockwise transformers for near-infinite context，2023】（称为SP-Ring），从两个不同的角度解决了SP固有的长期存在的内存和通信问题。对于这两种方法，每个计算设备都被分配了Q（查询）、K（键）、V（值）和O（输出）张量的不同部分，这些张量沿着序列维度被分割。它们在设备之间没有冗余存储，这是与早期SP设计【6】的主要区别。

SP-Ring机制

SP-Ring可以被看作是FlashAttention【【9】Flashattention: Fast and memory-efficient exact attention with io-awareness，2022】的分布式版本。如图2右侧所示，SP-Ring采用嵌套的两层循环，以分块方式协调通信和计算。在计算张量O分段的块时，如果所需的张量K和V块在本地不可用，则利用点对点（P2P）通信从其他设备获取它们。通信可以以环形（Ring）方式组织，其中每个设备同时发送和接收K、V块，从而使通信与计算重叠。


图2：SP-Ulysses与SP-Ring。

SP-Ulysses机制

SP-Ulysses利用All2All通信来处理Q、K、V和O张量的分段，如图2左侧所示。在All2All操作之后，这四个张量的分区从序列维度L变为注意力头数维度$h_c$。因此，每个注意力头的$softmax(QK^T)V$计算都完整地保留在单个设备上，并且可以使用底层的注意力操作库（如FlashAttention）来实现。

A2 方法细节

3. 统一的Ulysses-Ring序列并行

当前SP方法的局限性

目前，SP-Ulysses和SP-Ring都面临一些问题，这些问题阻碍了它们在实际应用中的有效性。

SP-Ulysses的限制

SP-Ulysses对注意力头的数量很敏感。DS-Ulysses的并行度不能超过注意力头的数量$h_c$。因此，它不适用于GQA（分组查询注意力）【【10】Gqa: Training generalized multi-query transformer models from multi-head checkpoints，2023】和MQA（多查询注意力）【【11】Fast transformer decoding: One write-head is all you need，2019】的场景。例如，Llama3-8B采用GQA，其KV头数量为8，这意味着使用DS-Ulysses SP时，最大SP度为8。然而，如果使用MQA且KV头数量为1，DS-Ulysses将无法工作。此外，由于张量并行也需要在$h_c$维度上进行划分，SP-Ulysses和TP存在冲突。

SP-Ring的限制

SP-Ring在计算和通信方面效率不高。Ring-Attention将Q、K、V、O张量分割成更小的块，这可能导致融合算子$Softmax(QK^T)V$的计算效率下降。即使通信和计算完全重叠，总执行时间也落后于DS-Ulysses。在使用因果掩码时，DS-Ring存在负载不均衡的问题。DS-Ring对$h_c$没有任何限制。

统一SP方法的提出

SP-Ulysses和SP-Ring目前被认为是SP的可选策略，选择其一则排斥另一。SP-Ring的作者在他们的ICLR公开评审答辩中强化了这一观点。目前，Megatron-DeepSpeed使用SP-Ulysses，而Megatron-LM则选择SP-Ring来实现其SP。然而，我们认为，它们可以作为一种统一的SP方法协同工作，而不是视其为竞争对手。

USP-Attention算法

如算法1所示，SP-Ring和SP-Ulysses以一种名为USP-Attention的混合并行方式组织起来，共同划分序列维度。SP进程组被分割成两个正交的进程组集合：一组SP-Ring进程组和一组SP-Ulysses进程组。为了更直观地理解，一个SP进程组可以被看作一个二维网格，SP-Ring在网格的每一列上运行，而SP-Ulysses在每一行上运行。例如，一个包含8个进程的进程组可以被看作2×4，其中一个大小为2的SP-Ulysses进程组ulysses_pg和一个大小为4的SP-Ring进程组ring_pg。这与数据并行和张量并行进程组的划分方式相同。USP-Attention的输入是沿序列维度划分后的Q、K、V张量的分片，输出是张量O的分片。张量分片的大小为$(bs, L/N, h_c, h_s)$。注意，当使用MQA时，K和V张量分片的$h_c$形状与Q张量的不同。在前向传播期间，scatter_idx设置为1，gather_idx设置为2。AllToAll4D合并L维度并划分张量Q、K、V的$h_c$维度，将它们转换为$(h_c/N_{ulysses}, bs, L, d)$。它们还将O张量沿L维度划分并合并$h_c$维度。在反向传播期间，scatter_idx设置为2，gather_idx设置为1。

SP-Ring的负载均衡问题

原始的SP-Ring在应用因果注意力时会引入负载不均衡问题，因为只需要计算$QK^T$的下三角矩阵。如图3所示，如果序列维度被均匀划分，计算任务在设备间分布不均。如图左侧所示，在4个GPU上，GPU3的计算负载几乎是GPU0的7倍。


图3：SP-Ring的负载均衡。

负载均衡解决方案

解决负载不均衡问题的方法是沿着序列维度对输入序列的token进行重排序，如图3右侧所示。在图中，输入序列包含16个token。在均匀分区下，GPU0处理token 0-3，而GPU3处理token 12-15。经过负载均衡分区的重排序操作后，GPU0现在处理token 0、1、14、15，而GPU3处理token 5、6、11、12。每个GPU处理的工作负载完美平衡，这是比条带注意力（striped attention）【【12】Striped attention: Faster ring attention for causal transformers，2023】更优的解决方案。

负载均衡算法实现

上述负载均衡方法相当直接，类似的实现已在Megatron-LM中应用。我们展示了其在统一SP方法中的应用。算法2在混合Ulysses和Ring并行中从全局输入序列中提取并重排序输入序列。对于仅涉及逐元素操作的位置编码，例如RoPE【【13】Roformer: Enhanced transformer with rotary position embedding，2024】，也需要对模型的位置编码参数应用相同的重排序操作。由于该操作应用于模型的输入token序列（一个长度为$bs \times L$的整数向量），因此负载均衡的额外开销可以忽略不计。

1: function LOCALBALANCELOCALSEQ(seq, ring_process_group, ulysses_process_group)
2:     ring_degree ← ring_process_group.get_world_size()
3:     ring_rank ← ring_process_group.get_rank()
4:     ulysses_rank ← ulysses_process_group.get_rank()
5:     seq_chunks ← seq.chunk(2 × ring_degree)
6:     reorder_seq ← concat([seq_chunks[r_rank], seq_chunks[2 × rd − r_rank − 1]])
7:     local_seq ← reorder_seq.chunk(ud)[u_rank]
8:     return local_seq
9: end function

统一SP的灵活性和鲁棒性

统一SP（Unified SP）非常灵活和鲁棒，允许Ulysses度和Ring度的各种组合，只要Ulysses度和Ring度的乘积等于SP度。当Ulysses度等于N时，它就变成SP-Ulysses；当Ring度等于N时，它就变成SP-Ring。因此，它涵盖了SP-Ulysses和SP-Ring的能力。

统一SP的优势

首先，统一SP可以消除SP-Ulysses的头数限制。例如，要在16度SP中运行具有$h_c=8$的llama3-8B，可以将Ulysses度设置为8，Ring度设置为2。其次，统一SP通过提供更鲁棒的通信模式，降低了对网络基础设施的带宽和拓扑要求。通过将Ulysses度设置为1到N之间的值，Attention通信模式将是P2P和All2All的混合，如图4所示。这种通信模式特别适合异构通信网络，允许All2All操作在高带宽互连中运行，而异步P2P通信在较低带宽部分运行。这适用于GPU通过PCIe Switch连接的节点内，或GPU节点之间通过以太网连接的集群等场景。

建议1：我们建议使用统一SP（Unified-SP）来替代SP-Ring和SP-Ulysses，因为它包含了两者的能力，同时提供了额外的好处。


图4：统一SP对网络硬件拓扑更具鲁棒性。

4. SP在4D并行中的应用

4D并行框架概述

作为一个新兴的并行方法，如何将SP整合到现有的数据并行（DP）、张量并行（TP）和流水线并行（PP）的混合并行框架中尚未被深入研究。本节将分析SP与DP、TP和PP之间的关系，并讨论设计包含序列维度并行的4D并行的最佳实践。

并行策略成本分析

如表2所示，我们分析了不同并行策略下标准Transformer模块的通信和内存成本。GQA的影响未在表中反映，但我们稍后会进行分析。“Communications Params”列中的条目表示在Transformer块的参数和梯度上进行的集体通信操作。它包括自注意力层中4个线性层以及FFN层中2个线性层的权重和偏置张量的参数/梯度，在GPT-2模型中总计为$12 \times O(d^2)$个元素。请注意，由于中间尺寸较小，llama3和llama2具有$9.37 \times O(d^2)$个元素。“Communications Act(ivation)”列中的条目表示在属于激活的单个隐藏状态张量上进行的集体通信，包含$bs \times L \times d$个元素。内存成本分为模型参数/梯度（P/G）、优化器状态（OS）和中间激活张量（Act）。

表2：标准Transformer块的SP、DP、TP和ZeRO的通信和内存成本比较

通信成本计算

通信成本代表带宽需求。它由通信元素数量乘以与集体通信算法相关的算法带宽（algobw）因子计算得出。对于AllReduce、AllGather、ReduceScatter和AllToAll的集体通信算法，各自的algobw因子分别为$2\frac{n-1}{n}$、$\frac{n-1}{n}$、$\frac{n-1}{n}$和1。在表中，为简化起见，我们将$O(\frac{n-1}{n})$项近似为$O(1)$。

内存成本计算

该表是为使用fp16(bf16)格式的混合精度训练构建的。模型参数和梯度的内存需求分别为P和G字节。优化器状态（OS），包括fp32格式的参数、动量和方差（在Adam优化器中），是fp16参数的6倍。峰值激活的内存需求为A字节。并行度为N。

SP vs. 数据并行 (DP)

在通信成本方面，SP不如DP。SP和DP在反向传播期间都需要对梯度进行allreduce操作。它们通信性能的差异在于注意力模块，SP为激活引入了额外的通信开销。当Ulysses度设置为大于1时，由于all2all操作，SP的通信开销将大于DP。当使用Ring方法时，尽管注意力的额外P2P通信被重叠，但它引入了额外的性能问题。理想性能仅能达到没有注意力通信的DP性能。在内存性能方面，SP和DP是等效的，因为它们都可以将激活占用减少到1/N。

建议2：我们建议如果可能，优先使用DP而不是SP。只有当批次大小（bs）不足以进行分区时，才应考虑是否采用SP。

SP vs. ZeRO

ZeRO【【14】Zero: Memory optimizations toward training trillion parameter models，2020】是一种分布式参数管理方法，通过在多个设备上分片优化器状态（ZeRO-1）、梯度（ZeRO-2）和参数（ZeRO-3）来减少每个计算设备的存储空间需求。优化器状态、梯度和参数的内存成本降低到原来的1/N。ZeRO也可以在SP进程组内工作，因为沿着批次维度（bs）或序列维度（L）进行分区与ZeRO的方法是等效的。ZeRO工作在大小为$N_{sp} \times N_{dp}$的统一进程组上，该进程组结合了SP和DP进程组。

SP vs. 张量并行 (TP) 与 TP-sp

由Megatron-LM【4】开创的张量并行（TP）方法，将模型参数分片到计算设备上。在TP中，并非所有激活张量都被分区并分布在多个计算设备上。因此，表2中反映的激活内存成本表示为$\alpha A$，其中$0 < \alpha < 1$。更精确的$\alpha$值请参考论文【【3】Reducing activation recomputation in large transformer models，2023】中的公式(2)。TP已被Megatron-LM序列并行【3】进一步改进，它将TP中的一个allreduce替换为一个allgather和一个reducescatter，因此将激活内存成本降低到A/P，代价是为注意力和FFN重做两次allgather。为了区分Megatron-LM序列并行与Ring和Ulysses序列并行，我们在此将其命名为TP-sp。

通信成本比较：SP vs. TP-sp

如表2所示，在通信成本方面，TP-sp高于SP-Ulysses和SP-Ring。首先，TP-sp的通信量大于SP-Ring，而后者的通信可以与计算重叠。此外，TP-sp的通信量不随并行度的增加而减少，而SP-Ulysses可以实现这一点。因此，在通信方面，TP-sp不如任何形式的序列并行。SP的激活通信成本较低，但需要同步梯度和参数。然而，与激活通信量相比，参数通信量较小，并且可以被计算重叠。GQA/MQA可以降低SP的通信成本，而TP-sp的通信成本保持不变。假设GQA组数为G，Ulysses和Ring对K和V的通信成本降低到1/G，激活通信成本分别降低到$4N \times O(bs \times L \times d) + 4N \times O(bs \times L \times d/G)$和$4 \times (bs \times L \times d/G)$。

建议4：我们认为，在大规模上，SP在通信成本方面优于TP-sp。GQA可以进一步降低SP的通信成本。

内存成本比较：SP vs. TP-sp

在内存成本方面，TP-sp优于SP。即使SP采用ZeRO-1/2来使OS和激活的内存成本与TP-sp对齐，参数的内存成本仍然更可观。然而，当SP采用ZeRO-3时，它实现了与TP-sp相似的内存成本。SP-Ulysses+ZeRO3正是DS-Ulysses【1】的作者用来扩展序列长度的策略。

可训练序列长度分析

基于以上分析，在相同的模型和输入配置下，简单地将并行策略从TP-sp切换到SP-Unified并不能增加可训练的序列长度；相反，很可能会导致序列长度的减少。然而，在高并行度下，SP仍然可以比TP-sp扩展序列长度。当序列非常长时，如前所述，参数通信量在总通信量中的比例相对较小。因此，ZeRO引入的额外allgather操作的通信开销影响有限。

建议5：我们建议将TP-sp切换到SP不能增加训练中的序列长度。SP+ZeRO3可以训练与TP-sp相似的序列长度。

扩展性优势

由于TP-sp的并行性存在固有限制，其上限为$h_c$，因此无法通过增加TP-sp并行度来进一步降低激活的内存成本。相比之下，SP可以通过利用SP-Ring技术继续扩展其并行度，从而能够在更大规模上训练更大的模型。

建议6：我们建议采用更高程度的SP并行（当头数受限时可能需要设置较大的ring度），以便在更多计算设备上训练长序列。这是TP-sp方法无法实现的优势。

SP vs. 流水线并行 (PP)

PP跨层划分Transformer块，因此它与在Transformer块内部划分张量的SP是互补的。因此，SP和PP完全兼容。由于SP可以与ZeRO形成统一的并行组进行集体通信，我们认为TP仍应放在4D并行组划分的最低维度。

建议7：我们建议在4D混合并行中，进程组维度的顺序从低到高是TP、SP-Ulysses、SP-Ring、ZeRO-DP、PP。

A4 实验环境

  • 硬件配置:
    • 集群1: L20 PCIe GPU集群 (8x L20)。
    • 集群2: 单个8xA100-SXM4 NVLink节点。
    • 集群3: 两个GPU节点，每个节点配备8块A800 80GB GPU，GPU间通过400GB/s NVLink连接，节点间通过1.6 Tbps RDMA连接。
  • 软件配置:
    • 代码库: 基于Megatron-LM（commit 2196398, 2024年4月12日）。
    • 实现: SP-Ring使用Megatron-LM原生的Context Parallel实现，SP-Ulysses基于作者的仓库代码开发。
    • 运行环境: Docker镜像 nvcr.io/nvidia/nemo:24.03。
    • 并行策略默认设置: 对数据并行（DP）和序列并行（SP）使用ZeRO-1。对张量并行（TP）始终应用序列并行优化（TP-sp）。不使用梯度累积或激活重计算。
  • 模型架构:
    • LLAMA2-7B
    • LLAMA3-8B (拥有8个注意力头)
  • 评测指标:
    • 吞吐量 (TFLOPS 或 iters/sec)。
    • 模型FLOPs利用率（MFU），经过调整以仅考虑因果掩码下的有效计算。

A4 实验结果

5.1 统一SP的性能

  • 实验内容: 在L20 PCIe GPU集群上评估SP-Unified在attention模块上的吞吐量（TFLOPS）。
  • 实验结果: 如表3所示，在8xL20配置下，对于32K和128K序列长度，当ulysses-degree=4且ring-degree=2时达到最佳吞吐量。此外，带有负载均衡的lb-ring性能优于标准Ring-Attention，且序列越长优势越明显。
  • 分析结论: 这个发现支持了建议1：SP-Unified非常适合异构网络。

表3：在8xL20 PCIe上SP-Unified的吞吐量（iters/sec）（仅前向）(Ring-Degree × Ulysses-Degree=8)

  • 实验内容: 在8xA100-SXM4 NVLink节点上重复attention基准测试。
  • 实验结果: 如表4所示，对于32K和128K序列长度，当ulysses-degree=8时（即纯SP-Ulysses）达到最高吞吐量，显著优于SP-Ring。
  • 分析结论: 这验证了第3节的论点，即SP-Ring虽然可以通过重叠隐藏通信开销，但会导致计算效率下降。

表4：在8xA100-SXM4 NVLink上SP-Unified的吞吐量（iters/sec）（仅前向）(Ring-Degree × Ulysses-Degree=8)

5.3 SP vs. DP

  • 实验内容: 在单节点A800 GPU上，使用LLAMA2-7B模型（全局批次大小为8）比较SP和DP的性能。
  • 实验结果: 如图5所示，在各种输入序列长度下，DP的性能均优于SP-Unified。
  • 分析结论: 证实了建议2的结论：应优先使用DP而非SP。


图5：在全局bs=8的LLAMA2-7B任务上比较DP和SP。

5.4 混合SP和TP

  • 实验内容: 在单节点8xA800上使用LLAMA2-7B进行混合并行性能测试。
  • 实验结果:
    1. 如表5所示，可实现的最长序列长度为64K，此时最佳性能由tp-degree=4和ulysses-degree=2的组合获得，比纯TP-sp快10%。当seqlen=64K时，纯SP会导致OOM。
    2. 当global-bs=16和seqlen=30K时，纯SP-Ulysses（ulysses-degree=8）性能最佳，吞吐量比纯TP-sp高26%。
  • 分析结论:
    1. 证实了建议5：SP的内存效率低于TP-sp。
    2. 在NVLINK连接的环境中，SP-Ulysses的通信效率在实践中可能优于TP-sp。

表5：在单个8xA800 NVLink节点上使用TP和SP-Unified的混合并行LLAMA2-7B的每GPU FLOPS。

  • 实验内容: 在两个节点（16x A800）上使用TP和SP-Unified的混合并行策略对LLAMA3-8B进行训练吞吐量基准测试。全局批次大小固定为1以最大化序列长度。
  • 实验结果:
    1. 64K和80K序列长度（表6）: 最佳性能由纯SP（ulysses-degree=4, ring-degree=4）实现，分别比SP-Ring快13%和12%。
    2. 内存限制: 纯SP在global-bs=2时会OOM，而TP+SP可以运行，但性能不一定更优。
    3. 120K序列长度（表6）: 最佳性能由tp-degree=4和ulysses-degree=2的混合策略实现，达到152.51 TFLOPS（MFU 0.49）。此时纯SP会OOM。
  • 分析结论:
    1. 再次证实建议1：统一SP在某些场景下优于单一SP策略。
    2. 再次证实建议5：SP的内存效率低于TP-sp，尤其是在处理极长序列时。

表6：在两个RDMA连接的8xA800 NVLink节点上使用TP和SP-Unified的混合并行LLAMA3-8B的每GPU FLOPS。

  • 实验内容: 在两个节点上探索LLAMA3-8B可训练序列长度的上限。
  • 实验结果: 如表7所示，最长序列长度（208K）由tp-degree=8和ring-degree=2的配置实现，MFU为0.47。此时纯SP因OOM无法运行。
  • 分析结论: 为了训练最长的序列，由于TP-sp内存效率更高，应将受注意力头数限制的并行度优先分配给TP-sp。

表7：在2个8xA800 NVLink节点上使用TP和SP-Unified探索LLAMA3-8B的序列长度上限。

5.5 收敛性

  • 实验内容: 在4个GPU上比较USP和DP的收敛性，测试了10K次迭代的损失曲线。
  • 实验结果: 如图6所示，USP和DP的损失曲线完全重合。
  • 分析结论: 验证了对RoPE和SP-Ring的负载均衡修改是正确的。


图6：在全局bs为4的LLAMA2-7B上比较DP和USP的损失。

A7 补充细节：未来工作

大规模集群上的SP

SP的潜力：我们相信SP对于超大规模LLM训练任务非常有益。目前，公开的大规模模型训练任务（超过10K GPU）【【17】Megascale: Scaling large language model training to more than 10,000 gpus，2024；【18】Introducing meta llama 3: The most capable openly available llm to date，2024】尚未使用SP，因为这些训练任务在2023年11月SP方法成熟之前就已启动。SP可以带来两个主要好处：首先，SP可以从序列长度维度引入分区，缓解批次大小限制带来的约束。例如，MegaScale项目【18】使用大的全局批次大小来增加DP度，这会影响收敛。我们可以增加SP度而不是DP度来减少全局批次大小，从而避免收敛问题。其次，增加SP度可以降低激活成本，允许在训练中采用更长的模型上下文长度。理论上，SP-Ring的并行度可以任意增加，而TP度受头数限制。我们在两个节点上的实验结果尚未展示出SP的这一优势。

SP+ZeRO-3

必要性：Megatron-LM目前不官方支持ZeRO-3，可能是因为TP-sp已经降低了参数和梯度的内存成本，这也是ZeRO-3的目标。我们已经声称ZeRO-3与SP高度兼容。鉴于Megatron-LM已经实现了SP-Ring，ZeRO-3成为一个必要的功能。

SP+MoE

兼容性：随着越来越多的模型转向专家混合（MoE）架构，研究SP与专家并行（EP）的结合变得日益重要。注意力模块上使用的序列并行与FFN模块是解耦的。因此，只要精心设计Attention和FFN模块之间的All2All通信，混合序列并行也可以与MoE兼容。

A5 结论

本文提出了一种统一DeepSpeed-Ulysses和Ring-Attention的序列并行方法。该方法不仅涵盖了这两种技术的能力，拓宽了适用范围，并在某些情况下提供了更优的性能。我们系统地分析了序列并行与其他已建立的并行方法之间的相互作用，并推导出一套最佳实践建议。这些建议通过在两个GPU节点上获得的实验结果得到了验证。

致谢

我们对来自腾讯的Zilin Zhu表示感谢。我们研究中使用的代码是基于他的GitHub仓库构建的，他也拥有本文图3的作者身份。