序列并行-负载均衡

序列并行

[🚀 序列并行 (Sequence Parallelism) 详解与负载均衡优化](#🚀 序列并行 (Sequence Parallelism) 详解与负载均衡优化)
- [1. 为什么需要序列并行？(Background)](#1. 为什么需要序列并行？(Background))
- [2. 序列并行是怎么切的？(Mechanism)](#2. 序列并行是怎么切的？(Mechanism))
- - 流派一：Megatron-SP (侧重省显存)
  - [流派二：Context Parallelism (DeepSpeed Ulysses / Ring Attention) (侧重长文本计算)](#流派二：Context Parallelism (DeepSpeed Ulysses / Ring Attention) (侧重长文本计算))
- [3. 核心痛点：负载不均衡 (The Load Imbalance)](#3. 核心痛点：负载不均衡 (The Load Imbalance))
- [4. 负载均衡优化方案 (Optimization Solutions)](#4. 负载均衡优化方案 (Optimization Solutions))
- - [方案 A：锯齿状切分 (Striped / Cyclic Partitioning) ------ **最推荐**](#方案 A：锯齿状切分 (Striped / Cyclic Partitioning) —— 最推荐)
  - [方案 B：动态调度与工作窃取 (Work Stealing)](#方案 B：动态调度与工作窃取 (Work Stealing))
  - [方案 C：2D 负载均衡 (Block-wise Ring Attention)](#方案 C：2D 负载均衡 (Block-wise Ring Attention))
- [5. 总结](#5. 总结)
- - [💡 总结图谱](#💡 总结图谱)

序列并行 (Sequence Parallelism, SP) && 负载均衡 (Load Balancing)

序列并行 主要是为了解决 Tensor Parallel (TP) 无法处理 超长上下文 (Long Context) 的问题（显存爆炸）。

🚀 序列并行 (Sequence Parallelism) 详解与负载均衡优化

1. 为什么需要序列并行？(Background)

在传统的 Megatron-LM Tensor Parallel (TP) 中：

切分维度： 主要切分 Hidden Dimension ( H H H)。
显存瓶颈： 虽然权重切分了，但 LayerNorm 和 Dropout 操作通常是复制 (Replicated) 的。这意味着每张卡都必须存储完整的 [ B a t c h , S e q L e n , H i d d e n ] [Batch, SeqLen, Hidden] [Batch,SeqLen,Hidden] 的激活值。
痛点： 当 Sequence Length ( L L L) 达到 128k 或 1M 时，光是存储这个激活值，单卡显存就爆了。

序列并行 (SP) 的核心思想：

不仅切分 Hidden 维度，在 LayerNorm 和 Attention 的维度上，沿着 Sequence Length ( L L L) 进行物理切分 。让单卡只需要存储 1 / N 1/N 1/N 的序列数据。

2. 序列并行是怎么切的？(Mechanism)

目前业界主要有两种 SP 的流派

流派一：Megatron-SP (侧重省显存)

这是 Megatron-LM v3 提出的。

原理： 在 Transformer 的 LayerNorm 和 Dropout 阶段，把数据按 L L L 切分。
通信变换： 将 TP 中的一次 All-Reduce 拆解为 Reduce-Scatter 和 All-Gather。
- Forward: MLP 输出 (TP行切) → Reduce-Scatter \xrightarrow{\text{Reduce-Scatter}} Reduce-Scatter 每个卡拿到 L / N L/N L/N 的数据 → LayerNorm \xrightarrow{\text{LayerNorm}} LayerNorm → All-Gather \xrightarrow{\text{All-Gather}} All-Gather 恢复全量 L L L 给 Attention。
收益： 激活值显存占用降低到 1 / N 1/N 1/N。

流派二：Context Parallelism (DeepSpeed Ulysses / Ring Attention) (侧重长文本计算)

这是目前处理 100k+ 长文本的主流（）。它直接把 Attention 计算也按 L L L 切分了。

KV Cache 切分：
- TP: 切 Head。
- SP: 切 Sequence。GPU 0 存第 0~1000 个 Token 的 KV，GPU 1 存第 1001~2000 个 Token 的 KV。
Q 切分：
- GPU 0 负责计算 Query 0~1000 的 Attention 结果。
- GPU 1 负责计算 Query 1001~2000 的 Attention 结果。

3. 核心痛点：负载不均衡 (The Load Imbalance)

在做 Context Parallelism 时，如果简单地把 Sequence 均分给不同的 GPU，会遇到严重的负载不均。

原因：Causal Mask (因果掩码 / 三角形计算)

Transformer Decoder 是自回归的，第 i i i 个 Token 只能看前 i i i 个 Token。Attention 矩阵是一个下三角矩阵。

场景： 序列长度 L = 8000 L=8000 L=8000，2 张卡 (GPU 0, GPU 1)。
朴素切分 (Naive Split)：
- GPU 0 (负责 Seq 0~3999):
  - 第 0 个 Token：看 1 个 KV。
  - 第 3999 个 Token：看 4000 个 KV。
  - 平均计算量： ∝ 1 2 × 4000 2 \propto \frac{1}{2} \times 4000^2 ∝21×40002 (梯形面积小)。
- GPU 1 (负责 Seq 4000~7999):
  - 第 4000 个 Token：看 4001 个 KV (前卡的 4000 + 自己的 1)。
  - 第 7999 个 Token：看 8000 个 KV。
  - 平均计算量： ∝ 1 2 × 8000 2 − GPU 0 Area \propto \frac{1}{2} \times 8000^2 - \text{GPU 0 Area} ∝21×80002−GPU 0 Area (梯形面积极大)。

后果： GPU 1 累死（计算量是 GPU 0 的 3 倍），GPU 0 算完后在那空转（Bubble）。木桶效应导致整体性能极差。

4. 负载均衡优化方案 (Optimization Solutions)

针对片内分布式或多卡集群，如何解决"三角形计算"带来的不均？

方案 A：锯齿状切分 (Striped / Cyclic Partitioning) ------ 最推荐

这是 DeepSpeed Ulysses 和 Ring Attention 常用的优化策略。

思路： 不要"切大块"，而是"切碎了轮流分"。
做法： 假设有 2 张卡。
- GPU 0 负责： Token { 0 , 2 , 4 , 6 , ... , 2 k } \{0, 2, 4, 6, \dots, 2k\} {0,2,4,6,...,2k}
- GPU 1 负责： Token { 1 , 3 , 5 , 7 , ... , 2 k + 1 } \{1, 3, 5, 7, \dots, 2k+1\} {1,3,5,7,...,2k+1}
效果：
- GPU 0 算的 Token 位置虽然靠前，但和 GPU 1 几乎是交错的。
- Token 2000 (GPU 0) 和 Token 2001 (GPU 1) 的计算负载（需要看的 KV 长度）几乎一样。
- 宏观上，两张卡的计算负载完全平衡。

方案 B：动态调度与工作窃取 (Work Stealing)

思路： 把 Attention 计算拆解成更小的 Tile (任务块) 。
- 任务池： T a s k ( Q i , K V a l l ) Task(Q_i, KV_{all}) Task(Qi,KVall)。
- 每个 Tile 的计算量是不同的（越靠后的 Q Q Q，计算量越大）。
做法：
- 使用一个全局调度器（或原子计数器）。
- 空闲的核心去领任务。
- 优化技巧： Longest Job First (LJF)。优先把靠后的、计算量大的 Token 分发出去，防止最后时刻出现长尾延迟。

方案 C：2D 负载均衡 (Block-wise Ring Attention)

思路： 将 Attention 矩阵的计算视为 ( Q , K ) (Q, K) (Q,K) 网格。
做法：
- 让 GPU 0 计算右上角的块（本来是被 Mask 掉的，可以不存数据，但可以分配逻辑任务）。
- 注：这种方法比较理论，实际工程中主要用方案 A。

5. 总结

序列并行有什么问题？怎么解决负载不均？

"序列并行（Context Parallelism）最大的挑战在于 Causal Mask 带来的计算不均衡。

现象： Attention 矩阵是下三角的。如果按顺序物理切分 Sequence，负责后端序列的计算单元（Core/GPU）计算量会远大于负责前端序列的单元（因为后面的 Token 要看所有的历史 KV），导致严重的 Pipeline Bubble。

通用解法 (Striped Partitioning)：

目前业界（如 DeepSpeed Ulysses）通用的解法是采用锯齿状（Cyclic）切分 。

不把序列连续切分，而是按 Token_ID % Device_Num 进行分配。

这样，每个 Device 都均匀地持有'头部'（计算量小）和'尾部'（计算量大）的 Token，从统计学上将计算负载完美拉平。

针对'片内分布式'的思考：

在片内多核架构下，这种 Striped 逻辑依然适用。

但考虑到片内 NoC (片上网络) 的带宽优势，我们可以做得更激进：

实现一个硬件感知的动态调度器 。将 Prefill 阶段的 Attention 计算拆解为细粒度的 Tile ，根据各个 Memory Block 的负载情况动态分发计算任务，这比静态的 Striped 切分更能适应动态变化的推理请求（尤其是 Chunked Prefill 场景下）。"

💡 总结图谱

方案	切分逻辑	负载均衡能力	通信复杂度	适用场景
朴素切分	`[0-N]` 给卡1, `[N-2N]` 给卡2	❌ 极差 (三角形不均)	低	BERT (双向注意力)
Striped切分	`0, 2, 4` 给卡1, `1, 3, 5` 给卡2	✅ 完美均衡	中 (需 All-to-All)	GPT/Llama (自回归)
Ring Attention	传递 KV Block 绕圈算	✅ 均衡 (时间换空间)	高 (重叠计算与通信)	超长文本 (1M+)