Intro & Background

• 大模型序列长度日益增加：GPT4o 128k，Claude3 200k，Gemini-1.5 Pro 2M

• 长序列训练的瓶颈：分布式的 Attention 计算很慢

  • Attention 计算复杂度是 O(N^2），FA 利用 online softmax，把存储复杂度降为 O(N)。但序列长度过长之后，还是需要在序列长度维度上做切分
  • 理论上，Context Parallel 通信量 O(N)，计算量是 O(N^2) ，在 N 很大的时候，应该是可以掩盖住 O(N) 通信的。但 profiling 出来发现盖不住，还是 communication-bound 了，因为 CP ranks 数量太多了

  

• ByteScale：主要解决的问题是，如何提高 “由短序列拼成的长序列“ 的训练效率（sequence packing），提出了 Hybrid Data Parallism（HDP）

  • 其核心思想是把 DP 和 CP 融合，统一进行调度。比如，能够上图中把 256 CP ranks 减小为了 39 HDP ranks，让 computation 能够盖住 communication
  • 原始的 DP+DP 方案还会引发 Imbalanced Computation，导致快的 DP 组必须等待慢的 DP 组。其根本原因是：虽然都是 1M 序列，但 1k 个 1k 下三角计算量之和，小于 1个 1M 下三角的 attention 计算量

序列维度上的并行方案

• TP + SP（Sequence-Parallel）

  • 来源于 Megatron-3 paper：Reducing Activation Recomputation in Large Transformer Models。不算是严格意义上的序列并行，只能算是序列切分
  • 把 dropout 和 LayerNorm 的激活沿序列维度切分了，切分数量等于张量并行数 tp，从而减少了激活值的内存占用
  • TP：正向 f=NULL，g=all-reduce；反向 f=all-reduce，g=NULL
  • TP+SP：正向 f=AG，g=RS；反向 f=RS，g=AG

  

  TP 下的 MLP 和 Attention 计算流

• DeepSpeed Ulysses：https://zhuanlan.zhihu.com/p/4496065391

  • Attention 层内是切 head，Attention 层外是切序列

    • QKV 在进 Attention 之前是切 sequence，然后做了 a2a，变为切 hidden，送进 Attention
    • Attention 之后，再通过 a2a 变回切 sequence

  • 优缺点：通信量小。但序列并行在 attention 里面是和 tp 一起去切 head，对 GQA 不友好

    

• Colossal AI 的 Ring Attention 方案：https://zhuanlan.zhihu.com/p/689067888

  • 在 attention 之外，是实打实的切分了序列
  • 在 attention 内，按序列切分了 KV，利用 ring-self attention（RSA）的思想：

    • 每个节点接收临近的 $k_i$，然后计算 $qk_i^T$，然后 allgather 得到 $qk^T$

    • 每个节点接收临近的 $v_i$，然后计算 $f(qk^T)v_i$，然后 allgather 得到 output

    • 优缺点：不受 head 数限制。通信量更大

      

• Nvidia mCore Context Parallel（CP）：https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/context_parallel.html

  从原理上基本等于 SP + CP（在 Ring-Self-Attention 基础上有一些额外优化）

  • 黄色的 AG/RS 是 SP 引入的
  • 深蓝的 AG/RS 是 CP 引入的

    • 前向的 AG 是为了收集完整的 KV 序列，对应其反向的 RS

    • 图中反向的额外 AG （Attn ouput 之前那个）是用于从切分后的 KV 做重计算，为了减少激活内存占用

      

      AG/RS 在唤醒拓扑下，可以用点对点通信实现

Motivation

1）Imbalanced Computation

通常情况下，8k/32k 等训练长度，一般都是多个样本（sample） 拼在一起得到一个序列（sequence），这个过程叫 sequence packing

但由于 Attention 的 O(n^2) 计算量，如果 packing 到一个序列中的 sample 有长有短， 那整个序列的计算时间其实会浮动的。例如，由 2 个 16k 拼成的 32k，计算时间会比 32 个 1k 拼成的 32k 序列计算时间更长。

这可能会导致下图的 dp bubble（快的 dp 组要等慢的 dp 组）。其中可以看到，第一个 dp 组的 pp bubble 明显也比第二个 dp 组的更长

2）Redundant Communication