第 7 课：序列并行 SP

对应官方章节：Sequence parallelism
前置：第 6 课（TP）
状态：未开始

本课目标

完成后你需要能够：

说明 SP 解决的问题：LayerNorm/dropout 需要完整隐藏维度，抵消 TP 的激活节省。
画出一个 transformer block 在 TP 与 TP+SP 下激活形状 (b, s, h) 的转换图。
解释 f/f* 与 g/g* 两组共轭操作在 forward/backward 中的角色。
说明 TP+SP 的通信量与纯 TP 相同的原因。

核心概念

1. 术语澄清：Megatron-SP 与 Context Parallelism (CP) 的本质区别

在大模型分布式系统领域，“序列并行（Sequence Parallelism）”这个词在历史上有两种完全不同的含义，面试与系统设计中极易混淆：

Megatron-SP（Korthikanti et al., 2022，本课核心）：专指与机内张量并行（TP）深度绑定的优化技术。它并不改变注意力机制的计算拓扑，而是专门针对 Transformer 中未被 TP 覆盖的区域（如 LayerNorm、RMSNorm、Dropout、残差连接等），通过沿序列维度（Sequence Dimension）切分来消除其显存冗余。
Context Parallelism（CP / 环形注意力 Ring-Attention，第 8 课）：专指跨物理节点将全模型的长文本序列（如 32k、128k、1M 上下文）整体切分并分布式流转计算的独立并行维度。

在讨论序列并行时，必须首先明确其上下文语境。

2. 传统 TP 的显存软肋：非 TP 区域的激活冗余

回顾第 6 课，张量并行（TP）成功将 MLP 和 Attention 核心的 GEMM 权重切分到各个 GPU 上，其内部中间激活张量的隐藏维度也被切分为 $\frac{h}{\text{TP}}$ 。

然而，如图 7-1 与图 7-2 的对比所示，在 GEMM 算子之外的非 TP 区域，系统面临严峻的显存回弹：

两种序列并行

Megatron SP（配合 TP）

LayerNorm、Dropout 等非 TP 区沿序列切分： $(B, \frac{S}{t}, d)$
进入 TP 区前 All-Gather，离开时 Reduce-Scatter，代替原来的 AllReduce
通信量与纯 TP 相同，非 TP 区激活降为 $1/t$

DeepSpeed Ulysses

输入按序列切分 $(\frac{S}{P}, H, d)$
All-to-All 转成按头切分 $(S, \frac{H}{P}, d)$ ，本地算注意力，再转回
要求头数（GQA 下为 KV 头数）能被 $P$ 整除

Megatron SP 主要省激活显存；Ulysses 每卡每次 All-to-All 传输约 $\frac{S d}{P}$ 量级的数据，随 $P$ 增大而减少，可直接复用 FlashAttention 内核

图 7-1：Megatron SP 与 DeepSpeed Ulysses

非张量并行区显存对比 (LN / Dropout / 残差)

标准 TP: 显存不切分
每卡持有全部 s 个 token
显存开销 = 100%

序列并行 SP: 沿序列切分
每卡仅持有 s / TP 个 token
显存开销 = 1 / TP

图 7-2：序列并行（SP）切分非张量并行区激活值显存分布对比

LayerNorm / RMSNorm 的归一化计算： $\mu = \frac{1}{h} \sum_{i=1}^h x_i, \quad \sigma^2 = \frac{1}{h} \sum_{i=1}^h (x_i - \mu)^2, \quad \text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta$ 均值 $\mu$ 与方差 $\sigma^2$ 是沿着整个隐藏特征维度 $h$ 进行统计约简的。若沿 $h$ 维度切分输入，各卡仅持有一部分通道特征，根本无法独立计算出正确的均值和方差，除非在 LayerNorm 内部强行插入跨卡通信。
全量激活冗余驻留：为了避免在每个 Norm 算子内发生低效的跨卡通信，原生 Megatron-LM 规定：在进入 LayerNorm、Dropout 以及残差连接之前，必须持有全量的特征张量 $(b, s, h)$ 。
后果：这意味着每个 TP 副本的 GPU 显存内都完整保存了一份完全相同的 $(b, s, h)$ 激活张量。随着模型层数 $L$ 与上下文长度 $s$ 的增长，这些未被切分的“公共激活值”占据了反向保留激活（Stored Activations）总量的 30%~50%，严重抵消了 TP 带来的显存收益。

3. 序列切分与通信共轭对（ $f/f^$ 与 $g/g^$ 原语闭环）

Megatron-SP 捕捉到了一个核心数学事实：LayerNorm、RMSNorm、GELU 和 Dropout 均是逐 Token（Token-wise）独立执行的算子 它们在序列维度 $s$ 上没有任何交叉依赖——第 1 个 Token 的归一化结果与第 100 个 Token 的归一化完全互不影响。

因此，完全可以沿序列维度 $s$ 将非 TP 区域的张量均分为 $\frac{s}{\text{TP}}$ 块，让每张卡仅驻留 $(b, \frac{s}{\text{TP}}, h)$ 的局部序列激活。

为了在 SP 区域（序列切分）与 TP 区域（隐藏维切分）之间无缝转换，Megatron 引入了优雅的共轭算子对（Conjugate Operator Pairs），其执行拓扑如图 7-1 与图 7-3 所示：

序列并行区 [s/tp, b, h]

AllGather (正向) / ReduceScatter (反向)

张量并行区 [s, b, h]

ReduceScatter (正向) / AllGather (反向)

1动态拼合完整序列进入 Column Linear

2完整序列执行局部矩阵乘与多头注意力

3Row Linear 局部输出直接规约切分序列

4规约后仅保留属于本卡的序列切片

图 7-3：Megatron-SP 巧妙利用 Ring 通信原语消除冗余通信的完整闭环

A. $g$ 与 $g^*$ 算子对（SP 与 TP 区域的转换核心）

前向转入 TP 区域（ $g$ 算子）：
- 在进入 Attention / MLP 的列并行矩阵乘之前，算子需要完整的 Token 语义（例如 Attention QKV 投影需与整体位置编码配合）。
- 系统发起一次 All-Gather(沿序列维度 $s$ )： $(b, \frac{s}{\text{TP}}, h) \xrightarrow{g = \text{All-Gather}(s)} (b, s, h)$
- 随后直接送入列并行线性层，输出沿隐藏维切分的张量 $(b, s, \frac{h}{\text{TP}})$ 。
前向转出 TP 区域（ $g^*$ 算子）：
- 行并行线性层计算后，各卡产生了部分和张量 $(b, s, h)$ 。
- 在原生 TP 中，此处必须执行一次 AllReduce 沿 $h$ 恢复完整和。
- 而在 SP 中，此处直接替换为 Reduce-Scatter(沿序列维度 $s$ )： $(b, s, \frac{h}{\text{TP}}) \xrightarrow{\text{Row Linear}} (b, s, h) \xrightarrow{g^* = \text{Reduce-Scatter}(s)} (b, \frac{s}{\text{TP}}, h)$
- 一石二鸟：Reduce-Scatter 既完成了行并行部分和的累加求和，又顺手将输出张量切分回了 $(b, \frac{s}{\text{TP}}, h)$ ，直接喂给后续的 LayerNorm 与 Dropout。

B. 共轭算子的反向行为与对偶性

在自动微分（Autograd）中，反向求导是前向有向无环图的伴随转置：

前向执行 $g$ （All-Gather 沿 $s$ ），反向求导对应的梯度汇聚算子正是其转置对偶：Reduce-Scatter 沿 $s$ 。
前向执行 $g^*$ （Reduce-Scatter 沿 $s$ ），反向求导对应的梯度恢复算子正是其转置对偶：All-Gather 沿 $s$ 。
这种代数上的高度自洽保证了前向与反向在数学上严格等价，完全无须破坏自动求导链条。

4. 通信量守恒定理：为什么 SP 几乎是“免费的午餐”？

直觉上，引入更多的通信算子会导致通信耗时增加。但严格的集合通信带宽分析证明：Megatron-SP 没有增加网络传输字节量。

回顾第 3 课的核心结论：在 Ring 环形通信算法下，一次 AllReduce 在底层完全等价于一次 Reduce-Scatter 紧接一次 All-Gather： $\text{AllReduce}(X) \equiv \text{All-Gather}(\text{Reduce-Scatter}(X))$ 其单步通信字节数关系为： $V_{\text{AllReduce}} = 2 \times \frac{\text{TP} - 1}{\text{TP}} \times \text{Size}, \quad V_{\text{Reduce-Scatter}} = V_{\text{All-Gather}} = \frac{\text{TP} - 1}{\text{TP}} \times \text{Size}$

纯 TP 模式每 Block 通信量：前向包含 2 次 AllReduce（Attention 输出 + MLP 输出）： $V_{\text{TP}} = 2 \times V_{\text{AllReduce}} = 4 \times \frac{\text{TP}-1}{\text{TP}} \times (b \cdot s \cdot h \cdot 2)$
TP + SP 模式每 Block 通信量：前向包含 2 次 All-Gather（进入 Col Linear）+ 2 次 Reduce-Scatter（退出 Row Linear）： $V_{\text{TP+SP}} = 2 \times (V_{\text{All-Gather}} + V_{\text{Reduce-Scatter}}) = 4 \times \frac{\text{TP}-1}{\text{TP}} \times (b \cdot s \cdot h \cdot 2)$
结论： $V_{\text{TP+SP}} = V_{\text{TP}}$ 。操作次数从 2 次变为 4 次，每次数据量减半；在机内 NVLink 上多出的两次启动延迟（约 $2\mu s$ ）很小。
- 作为回报，模型所有非 TP 区域的驻留激活显存被直接消减了 $\text{TP}$ 倍（例如 8 卡下直接砍掉 87.5% 激活显存）。

具体演示与数值推演

以一个 70B 模型（隐藏层维度 $h=8192$ ，层数 $L=80$ ），在 $\text{TP}=8$ 、序列长度 $s=16384$ （16k 上下文）、单卡批次 $\text{MBS}=1$ （BF16 精度，每元素 2 字节）下的单层驻留激活为例：

未开启 SP（原生 TP-only）：按 Korthikanti 等的激活公式，TP 下每层激活为 $sbh\left(10 + \frac{24}{t} + \frac{5as}{ht}\right)$ 字节，其中未被 TP 切分的部分约 $10sbh$ ：两个 LayerNorm 的输入（ $4sbh$ ）、QKV 投影与 MLP 第一层的输入（ $4sbh$ ）、两个 Dropout 掩码（ $2sbh$ ，每元素 1 字节）。单层为： $M_{\text{act/layer}} \approx 10 \times 16384 \times 1 \times 8192\text{ 字节} \approx 1.34\text{ GB}$ 全网 80 层累加： $M_{\text{total\_act}} \approx 80 \times 1.34\text{ GB} \approx 107.4\text{ GB}$ 仅这部分未切分的激活就已超过单卡 80 GB，不开 SP 时只能依赖重计算。
开启 SP（TP + SP 协同）：这些张量沿序列切为 $\frac{16384}{8} = 2048$ 个 token，单层降为： $M_{\text{act/layer\_SP}} \approx 10 \times 2048 \times 1 \times 8192\text{ 字节} \approx 167.8\text{ MB} \quad (\text{为原来的 } 1/8)$ 80 层合计约 $13.4\text{ GB}$ ，省出约 $94\text{ GB}$ 。

代码填空题

模拟一个 transformer block 内激活形状 (s, h) 的变化，比较 TP-only 与 TP+SP 的峰值激活。

Code

def trace_block(
    seq: int, hidden: int, tp: int, enable_sp: bool, batch: int = 1,
) -> list[tuple[str, int, int]]:
    """
    返回 [(区域名, s, h)]：前向经过 block 时激活形状的轨迹（每 token 两字节 BF16 可算字节数）。

    算子表（对 (s, h) 的作用）：
      split_seq    : s -> s/tp      （SP 区入口 / g* 出口）
      gather_seq   : s -> s*tp      （g：SP -> TP）
      split_hidden : h -> h/tp      （col-linear 切 h）
      gather_hidden: h -> h*tp      （row-linear 恢复 h）
      noop         : 不变

    区域顺序：
      layernorm → col_linear → gelu → row_linear → dropout
    """
    def split_seq(s, h): return (s // tp, h)
    def gather_seq(s, h): return (s * tp, h)
    def split_hidden(s, h): return (s, h // tp)
    def gather_hidden(s, h): return (s, h * tp)
    def noop(s, h): return (s, h)

    region_ops = {
        "layernorm": ______,   # 填空：TP-only 下 LN 输入就是全量 (s,h)，用什么算子？
        "col_linear": ______,  # 填空：col-linear 沿隐藏维切分
        "gelu": ______,        # 填空：逐元素，形状不变
        "row_linear": ______,  # 填空：row-linear 恢复隐藏维
        "dropout": ______,     # 填空：TP-only 下 dropout 保持全量
    }
    sp_ops = {
        # 启用 SP 时各区域用的算子（提示：LN/dropout 在 SP 区沿序列切，
        # col/row 在 TP 区沿隐藏维切）
        "layernorm": ______,
        "col_linear": ______,
        "gelu": ______,
        "row_linear": ______,
        "dropout": ______,
    }
    plan = sp_ops if enable_sp else region_ops

    s, h = seq, hidden
    trace = []
    for name, op in plan.items():
        if enable_sp and name == "col_linear":
            # SP -> TP 过渡：g = all-gather 沿序列维，恢复完整 s
            s, h = gather_seq(s, h)
            trace.append(("g (all-gather seq)", s, h, False))   # 瞬态通信缓冲
        s, h = op(s, h)
        # SP 模式下 row-linear 的 (s, h) 输出是瞬态（随后被 g* 切回序列分片），
        # 反向真正需要的是它的输入（上一区域的 (s, h/tp)）——所以标记为不驻留。
        persists = not (enable_sp and name == "row_linear")
        trace.append((name, s, h, persists))                    # 驻留激活（要存给反向）
        if enable_sp and name == "row_linear":
            trace.append(("row_out (full h)", s, h, False))      # 瞬态：随后被 g* 切回
            s, h = split_seq(s, h)
            trace.append(("g* (reduce-scatter seq)", s, h, False))
    return trace

def max_activation_bytes(trace: list[tuple[str, int, int, bool]], batch: int) -> int:
    """
    峰值【驻留】激活 = max(需要保存给反向的激活字节数)，BF16 每元素 2 字节。
    g/g* 与 row_out 是瞬态全量张量（存在通信/临时缓冲区里），不计入驻留激活；
    但它们会出现在峰值显存里——面试时要主动区分"驻留激活"与"瞬态缓冲"。
    """
    return max(
        ______   # 填空：只统计驻留条目 (stored=True)，b*s*h*2
        for name, s, h, stored in trace
        if stored
    )

def bf16_bytes(s: int, h: int, batch: int) -> int:
    return batch * s * h * 2

if __name__ == "__main__":
    seq, hidden, tp, bs = 4096, 4096, 8, 1

    t_tp = trace_block(seq, hidden, tp, enable_sp=False, batch=bs)
    t_sp = trace_block(seq, hidden, tp, enable_sp=True, batch=bs)

    print("TP-only:")
    for name, s, h, stored in t_tp:
        tag = " " if stored else "*"
        print(f"  {tag}{name:22s} s={s:5d} h={h:5d}  {bf16_bytes(s, h, bs)/1e6:8.1f} MB")
    print("TP+SP（* = 瞬态全量，不计入驻留激活）:")
    for name, s, h, stored in t_sp:
        tag = " " if stored else "*"
        print(f"  {tag}{name:26s} s={s:5d} h={h:5d}  {bf16_bytes(s, h, bs)/1e6:8.1f} MB")

    m_tp = max_activation_bytes(t_tp, bs)
    m_sp = max_activation_bytes(t_sp, bs)
    print(f"\nTP-only 峰值 = {m_tp/1e6:.1f} MB = b·s·h·2 = {bs*seq*hidden*2/1e6:.1f} MB")
    print(f"TP+SP   峰值 = {m_sp/1e6:.1f} MB = b·s·h·2/tp = {bs*seq*hidden*2/tp/1e6:.1f} MB")
    assert m_sp * tp == m_tp
    print("验证：TP+SP 峰值激活 = TP-only 的 1/tp ✓")

三个问答题

Q1

为什么 LayerNorm 不能像 MLP 那样沿隐藏维切分？请写出 LayerNorm 公式并指出哪一步依赖完整 h。dropout 为什么也需要 SP（提示：随机掩码的模式）？

Q2

教材用 f/f* 描述 TP 区域的入口与出口（与非 TP 区域的边界）、用 g/g* 描述 SP 与 TP 区域的交界。请分别写出前向和反向下 f、f* 各是什么操作（no-op 还是 all-reduce），并解释为什么反向恰好互换。SP 区域为什么刻意不用 all-reduce？

Q3

教材说 TP+SP “通信操作数是 TP 的两倍（2 all-gather + 2 reduce-scatter vs 2 all-reduce），但总通信量等价”。请用第 3 课的结论（all-reduce = reduce-scatter + all-gather）证明这一点，并指出这为什么使 SP 几乎免费。

检查与通过标准

总分 10 分：代码正确 4 分（变换计划、g/g* 过渡、峰值公式与验证）、三题各 2 分、通过线 8 分。

一票否决项：

把 SP 与 CP 混为一谈（本书定义下）。
认为 SP 沿隐藏维切分 LayerNorm 的输入。
说不出 g* 是 reduce-scatter（把 row-linear 的求和与序列切分合并）或方向错误。
声称 SP 增加了通信量（而非通信次数增加、总量等价）。

💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。

Code

def trace_block(
    seq: int, hidden: int, tp: int, enable_sp: bool, batch: int = 1,
) -> list[tuple[str, int, int]]:
    """
    返回 [(区域名, s, h)]：前向经过 block 时激活形状的轨迹（每 token 两字节 BF16 可算字节数）。

    算子表（对 (s, h) 的作用）：
      split_seq    : s -> s/tp      （SP 区入口 / g* 出口）
      gather_seq   : s -> s*tp      （g：SP -> TP）
      split_hidden : h -> h/tp      （col-linear 切 h）
      gather_hidden: h -> h*tp      （row-linear 恢复 h）
      noop         : 不变

    区域顺序：
      layernorm → col_linear → gelu → row_linear → dropout
    """
    def split_seq(s, h): return (s // tp, h)
    def gather_seq(s, h): return (s * tp, h)
    def split_hidden(s, h): return (s, h // tp)
    def gather_hidden(s, h): return (s, h * tp)
    def noop(s, h): return (s, h)

    region_ops = {
        "layernorm": noop,   # 填空：TP-only 下 LN 输入就是全量 (s,h)，用什么算子？
        "col_linear": split_hidden,  # 填空：col-linear 沿隐藏维切分
        "gelu": noop,        # 填空：逐元素，形状不变
        "row_linear": gather_hidden,  # 填空：row-linear 恢复隐藏维
        "dropout": noop,     # 填空：TP-only 下 dropout 保持全量
    }
    sp_ops = {
        # 启用 SP 时各区域用的算子（提示：LN/dropout 在 SP 区沿序列切，
        # col/row 在 TP 区沿隐藏维切）
        "layernorm": split_seq,
        "col_linear": split_hidden,
        "gelu": noop,
        "row_linear": gather_hidden,
        "dropout": noop,
    }
    plan = sp_ops if enable_sp else region_ops

    s, h = seq, hidden
    trace = []
    for name, op in plan.items():
        if enable_sp and name == "col_linear":
            # SP -> TP 过渡：g = all-gather 沿序列维，恢复完整 s
            s, h = gather_seq(s, h)
            trace.append(("g (all-gather seq)", s, h, False))   # 瞬态通信缓冲
        s, h = op(s, h)
        # SP 模式下 row-linear 的 (s, h) 输出是瞬态（随后被 g* 切回序列分片），
        # 反向真正需要的是它的输入（上一区域的 (s, h/tp)）——所以标记为不驻留。
        persists = not (enable_sp and name == "row_linear")
        trace.append((name, s, h, persists))                    # 驻留激活（要存给反向）
        if enable_sp and name == "row_linear":
            trace.append(("row_out (full h)", s, h, False))      # 瞬态：随后被 g* 切回
            s, h = split_seq(s, h)
            trace.append(("g* (reduce-scatter seq)", s, h, False))
    return trace

def max_activation_bytes(trace: list[tuple[str, int, int, bool]], batch: int) -> int:
    """
    峰值【驻留】激活 = max(需要保存给反向的激活字节数)，BF16 每元素 2 字节。
    g/g* 与 row_out 是瞬态全量张量（存在通信/临时缓冲区里），不计入驻留激活；
    但它们会出现在峰值显存里——面试时要主动区分"驻留激活"与"瞬态缓冲"。
    """
    return max(
        max(batch * s * h * 2 for _, s, h, stored in trace if stored)   # 填空：只统计驻留条目 (stored=True)，b*s*h*2
        for name, s, h, stored in trace
        if stored
    )

def bf16_bytes(s: int, h: int, batch: int) -> int:
    return batch * s * h * 2

if __name__ == "__main__":
    seq, hidden, tp, bs = 4096, 4096, 8, 1

    t_tp = trace_block(seq, hidden, tp, enable_sp=False, batch=bs)
    t_sp = trace_block(seq, hidden, tp, enable_sp=True, batch=bs)

    print("TP-only:")
    for name, s, h, stored in t_tp:
        tag = " " if stored else "*"
        print(f"  {tag}{name:22s} s={s:5d} h={h:5d}  {bf16_bytes(s, h, bs)/1e6:8.1f} MB")
    print("TP+SP（* = 瞬态全量，不计入驻留激活）:")
    for name, s, h, stored in t_sp:
        tag = " " if stored else "*"
        print(f"  {tag}{name:26s} s={s:5d} h={h:5d}  {bf16_bytes(s, h, bs)/1e6:8.1f} MB")

    m_tp = max_activation_bytes(t_tp, bs)
    m_sp = max_activation_bytes(t_sp, bs)
    print(f"\nTP-only 峰值 = {m_tp/1e6:.1f} MB = b·s·h·2 = {bs*seq*hidden*2/1e6:.1f} MB")
    print(f"TP+SP   峰值 = {m_sp/1e6:.1f} MB = b·s·h·2/tp = {bs*seq*hidden*2/tp/1e6:.1f} MB")
    assert m_sp * tp == m_tp
    print("验证：TP+SP 峰值激活 = TP-only 的 1/tp ✓")

补全后的代码（关键填空处）

region_ops = {
        "layernorm": noop,             # TP-only：LN 输入本来就是全量 (s, h)
        "col_linear": split_hidden,    # col-linear 切 h
        "gelu": noop,                  # 逐元素
        "row_linear": gather_hidden,   # row-linear 恢复 h
        "dropout": noop,               # TP-only：全量
    }
    sp_ops = {
        "layernorm": split_seq,        # SP 区：沿序列切
        "col_linear": split_hidden,    # TP 区：沿隐藏维切
        "gelu": noop,
        "row_linear": gather_hidden,
        "dropout": noop,               # SP 区：保持序列切分
    }
    # 过渡：col_linear 前 g=gather_seq（SP→TP）；row_linear 后 g*=split_seq（TP→SP）
    # 驻留标记：g/g*/row_out 为瞬态（存在通信/临时缓冲区），不计入驻留激活

def max_activation_bytes(trace, batch):
    return max(batch * s * h * 2 for _, s, h, stored in trace if stored)

运行结果（seq=4096, h=4096, tp=8）：

TP-only：LN/row/dropout 处形状 (4096, 4096) → 驻留峰值 33.6 MB = b·s·h·2
TP+SP ：layernorm (512,4096)、col/gelu (4096,512)、dropout (512,4096)
        → 驻留峰值 4.2 MB = b·s·h·2/tp ✓（g、row_out、g* 是瞬态全量）
验证：TP+SP 驻留峰值 = TP-only 的 1/tp ✓

三个问答题要点

Q1

LayerNorm：μ 与 σ² 沿 h 维求均值/方差（LayerNorm(x) = γ·(x−μ)/√(σ²+ε)+β），每 rank 只有 h/tp 无法算出正确的统计量 → 必须全量 h；
dropout：掩码是随机模式，各 rank 必须用相同种子生成一致的掩码（否则激活模式不同步），且掩码要覆盖全量张量；
所以这些”便宜”操作反而是激活内存大户 → 沿序列维切分（SP）把它们也分片。

Q2

f/f*（TP 区域的入口/出口，不开 SP 时使用）：前向 f=no-op（输入已同步）、f*=all-reduce（保证正确性）；反向互换（梯度形状是前向的转置，谁前向是 no-op 谁反向就是 all-reduce）；
g/g*（SP 与 TP 交界）：前向 g=all-gather（序列拼回全量）、g*=reduce-scatter（合并 row-linear 的部分和 + 切回序列分片）；反向为共轭互换；
SP 区域不用 all-reduce 的原因：all-reduce 需要物化全量激活（先 gather 再 reduce），峰值内存回到 b·s·h，SP 的初衷就没了；用 reduce-scatter 把”求和”与”切分”一步完成。

Q3

由 all-reduce = reduce-scatter + all-gather，纯 TP 的 2 次 all-reduce 与 TP+SP 的 2 次 all-gather + 2 次 reduce-scatter 字节数相同，只是操作数翻倍（推导详见第 4 节）。
额外代价只有实现复杂度和更多次启动延迟，带宽成本与 TP 相同，所以说 SP 几乎免费。

通过要点

术语：本书 SP 专指 TP 的配套技术（LN/dropout 沿序列切）；独立的序列切分叫 CP（第 8 课）。
激活形状：TP 区域 h 切、SP 区域 s 切；驻留激活峰值 b·s·h/tp；g/g* 的瞬态全量张量只存在于通信/临时缓冲区（面试要主动区分）。

官方主参考

← 上一课
第 6 课：张量并行 TP 下一课 →
第 8 课：上下文并行 CP

--- title: "第 7 课：序列并行 SP" jupyter: python3 --- - 对应官方章节：Sequence parallelism - 前置：第 6 课（TP） - 状态：未开始 ## 本课目标完成后你需要能够： - 说明 SP 解决的问题：LayerNorm/dropout 需要完整隐藏维度，抵消 TP 的激活节省。 - 画出一个 transformer block 在 TP 与 TP+SP 下激活形状 `(b, s, h)` 的转换图。 - 解释 f/f\* 与 g/g\* 两组共轭操作在 forward/backward 中的角色。 - 说明 TP+SP 的通信量与纯 TP 相同的原因。 ## 核心概念 ### 1. 术语澄清：Megatron-SP 与 Context Parallelism (CP) 的本质区别在大模型分布式系统领域，“序列并行（Sequence Parallelism）”这个词在历史上有两种完全不同的含义，面试与系统设计中极易混淆： 1. **Megatron-SP（Korthikanti et al., 2022，本课核心）**：专指与**机内张量并行（TP）深度绑定**的优化技术。它并不改变注意力机制的计算拓扑，而是专门针对 Transformer 中未被 TP 覆盖的区域（如 LayerNorm、RMSNorm、Dropout、残差连接等），通过沿序列维度（Sequence Dimension）切分来消除其显存冗余。 2. **Context Parallelism（CP / 环形注意力 Ring-Attention，第 8 课）**：专指跨物理节点将全模型的长文本序列（如 32k、128k、1M 上下文）整体切分并分布式流转计算的独立并行维度。在讨论序列并行时，必须首先明确其上下文语境。 ### 2. 传统 TP 的显存软肋：非 TP 区域的激活冗余回顾第 6 课，张量并行（TP）成功将 MLP 和 Attention 核心的 GEMM 权重切分到各个 GPU 上，其内部中间激活张量的隐藏维度也被切分为 $\frac{h}{\text{TP}}$。然而，如**图 7-1 与图 7-2** 的对比所示，在 GEMM 算子之外的**非 TP 区域**，系统面临严峻的显存回弹： :::: {.vfig} ::: {.vfig-title} 两种序列并行 ::: ::::: {.vfig-grid style="--cols:2"} ::: {.vfig-card .c-blue} **Megatron SP（配合 TP）** - LayerNorm、Dropout 等非 TP 区沿序列切分：$(B, \frac{S}{t}, d)$ - 进入 TP 区前 All-Gather，离开时 Reduce-Scatter，代替原来的 AllReduce - 通信量与纯 TP 相同，非 TP 区激活降为 $1/t$ ::: ::: {.vfig-card .c-purple} **DeepSpeed Ulysses** - 输入按序列切分 $(\frac{S}{P}, H, d)$ - All-to-All 转成按头切分 $(S, \frac{H}{P}, d)$，本地算注意力，再转回 - 要求头数（GQA 下为 KV 头数）能被 $P$ 整除 ::: ::::: ::: {.vfig-key} Megatron SP 主要省激活显存；Ulysses 每卡每次 All-to-All 传输约 $\frac{S d}{P}$ 量级的数据，随 $P$ 增大而减少，可直接复用 FlashAttention 内核 ::: :::: <p class="caption">图 7-1：Megatron SP 与 DeepSpeed Ulysses</p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-lr"><div class="fd-sg" data-nid="NonTP" data-depth="0"><div class="fd-sg-title">非张量并行区显存对比 (LN / Dropout / 残差)</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="Standard">标准 TP: 显存不切分<br>每卡持有全部 s 个 token<br>显存开销 = 100%</div><div class="fd-node fd-box" data-nid="SP">序列并行 SP: 沿序列切分<br>每卡仅持有 s / TP 个 token<br>显存开销 = 1 / TP</div></div></div></div></div></div></div></div> ``` <p class="caption" align="center"><em>图 7-2：序列并行（SP）切分非张量并行区激活值显存分布对比</em></p> - **LayerNorm / RMSNorm 的归一化计算**： $$\mu = \frac{1}{h} \sum_{i=1}^h x_i, \quad \sigma^2 = \frac{1}{h} \sum_{i=1}^h (x_i - \mu)^2, \quad \text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta$$ 均值 $\mu$ 与方差 $\sigma^2$ 是沿着**整个隐藏特征维度 $h$** 进行统计约简的。若沿 $h$ 维度切分输入，各卡仅持有一部分通道特征，根本无法独立计算出正确的均值和方差，除非在 LayerNorm 内部强行插入跨卡通信。 - **全量激活冗余驻留**：为了避免在每个 Norm 算子内发生低效的跨卡通信，原生 Megatron-LM 规定：在进入 LayerNorm、Dropout 以及残差连接之前，必须持有全量的特征张量 $(b, s, h)$。 - **后果**：这意味着每个 TP 副本的 GPU 显存内都完整保存了一份完全相同的 $(b, s, h)$ 激活张量。随着模型层数 $L$ 与上下文长度 $s$ 的增长，这些未被切分的“公共激活值”占据了反向保留激活（Stored Activations）总量的 30%~50%，严重抵消了 TP 带来的显存收益。 ### 3. 序列切分与通信共轭对（$f/f^*$ 与 $g/g^*$ 原语闭环） Megatron-SP 捕捉到了一个核心数学事实：**LayerNorm、RMSNorm、GELU 和 Dropout 均是逐 Token（Token-wise）独立执行的算子** 它们在序列维度 $s$ 上没有任何交叉依赖——第 1 个 Token 的归一化结果与第 100 个 Token 的归一化完全互不影响。因此，完全可以**沿序列维度 $s$ 将非 TP 区域的张量均分为 $\frac{s}{\text{TP}}$ 块**，让每张卡仅驻留 $(b, \frac{s}{\text{TP}}, h)$ 的局部序列激活。为了在 SP 区域（序列切分）与 TP 区域（隐藏维切分）之间无缝转换，Megatron 引入了优雅的**共轭算子对（Conjugate Operator Pairs）**，其执行拓扑如**图 7-1 与图 7-3** 所示： ```{=html} <div class="fd fd-seq" role="img"><div class="sd-grid" style="--n:4"><div class="sd-life" style="grid-column:1;grid-row:1/6"></div><div class="sd-life" style="grid-column:2;grid-row:1/6"></div><div class="sd-life" style="grid-column:3;grid-row:1/6"></div><div class="sd-life" style="grid-column:4;grid-row:1/6"></div><div class="sd-actor" style="grid-column:1;grid-row:1">序列并行区 [s/tp, b, h]</div><div class="sd-actor" style="grid-column:2;grid-row:1">AllGather (正向) / ReduceScatter (反向)</div><div class="sd-actor" style="grid-column:3;grid-row:1">张量并行区 [s, b, h]</div><div class="sd-actor" style="grid-column:4;grid-row:1">ReduceScatter (正向) / AllGather (反向)</div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:2;--k:2"><div class="sd-text"><span class="sd-num">1</span>动态拼合完整序列进入 Column Linear</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:3;--k:2"><div class="sd-text"><span class="sd-num">2</span>完整序列执行局部矩阵乘与多头注意力</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:3/5;grid-row:4;--k:2"><div class="sd-text"><span class="sd-num">3</span>Row Linear 局部输出直接规约切分序列</div><div class="sd-line"></div></div><div class="sd-msg sd-l sd-h-arrow sd-dash" style="grid-column:1/5;grid-row:5;--k:4"><div class="sd-text"><span class="sd-num">4</span>规约后仅保留属于本卡的序列切片</div><div class="sd-line"></div></div></div></div> ``` <p class="caption" align="center"><em>图 7-3：Megatron-SP 巧妙利用 Ring 通信原语消除冗余通信的完整闭环</em></p> #### A. $g$ 与 $g^*$ 算子对（SP 与 TP 区域的转换核心） - **前向转入 TP 区域（$g$ 算子）**： - 在进入 Attention / MLP 的列并行矩阵乘之前，算子需要完整的 Token 语义（例如 Attention QKV 投影需与整体位置编码配合）。 - 系统发起一次 **All-Gather(沿序列维度 $s$)**： $$(b, \frac{s}{\text{TP}}, h) \xrightarrow{g = \text{All-Gather}(s)} (b, s, h)$$ - 随后直接送入列并行线性层，输出沿隐藏维切分的张量 $(b, s, \frac{h}{\text{TP}})$。 - **前向转出 TP 区域（$g^*$ 算子）**： - 行并行线性层计算后，各卡产生了部分和张量 $(b, s, h)$。 - 在原生 TP 中，此处必须执行一次 AllReduce 沿 $h$ 恢复完整和。 - **而在 SP 中，此处直接替换为 Reduce-Scatter(沿序列维度 $s$)**： $$(b, s, \frac{h}{\text{TP}}) \xrightarrow{\text{Row Linear}} (b, s, h) \xrightarrow{g^* = \text{Reduce-Scatter}(s)} (b, \frac{s}{\text{TP}}, h)$$ - **一石二鸟**：`Reduce-Scatter` 既完成了行并行部分和的累加求和，又顺手将输出张量切分回了 $(b, \frac{s}{\text{TP}}, h)$，直接喂给后续的 LayerNorm 与 Dropout。 #### B. 共轭算子的反向行为与对偶性在自动微分（Autograd）中，反向求导是前向有向无环图的伴随转置： - 前向执行 $g$（All-Gather 沿 $s$），反向求导对应的梯度汇聚算子正是其转置对偶：**Reduce-Scatter 沿 $s$**。 - 前向执行 $g^*$（Reduce-Scatter 沿 $s$），反向求导对应的梯度恢复算子正是其转置对偶：**All-Gather 沿 $s$**。 - 这种代数上的高度自洽保证了前向与反向在数学上严格等价，完全无须破坏自动求导链条。 ### 4. 通信量守恒定理：为什么 SP 几乎是“免费的午餐”？直觉上，引入更多的通信算子会导致通信耗时增加。但严格的集合通信带宽分析证明：**Megatron-SP 没有增加网络传输字节量**。回顾第 3 课的核心结论：在 Ring 环形通信算法下，一次 AllReduce 在底层完全等价于一次 Reduce-Scatter 紧接一次 All-Gather： $$\text{AllReduce}(X) \equiv \text{All-Gather}(\text{Reduce-Scatter}(X))$$ 其单步通信字节数关系为： $$V_{\text{AllReduce}} = 2 \times \frac{\text{TP} - 1}{\text{TP}} \times \text{Size}, \quad V_{\text{Reduce-Scatter}} = V_{\text{All-Gather}} = \frac{\text{TP} - 1}{\text{TP}} \times \text{Size}$$ - **纯 TP 模式每 Block 通信量**：前向包含 2 次 AllReduce（Attention 输出 + MLP 输出）： $$V_{\text{TP}} = 2 \times V_{\text{AllReduce}} = 4 \times \frac{\text{TP}-1}{\text{TP}} \times (b \cdot s \cdot h \cdot 2)$$ - **TP + SP 模式每 Block 通信量**：前向包含 2 次 All-Gather（进入 Col Linear）+ 2 次 Reduce-Scatter（退出 Row Linear）： $$V_{\text{TP+SP}} = 2 \times (V_{\text{All-Gather}} + V_{\text{Reduce-Scatter}}) = 4 \times \frac{\text{TP}-1}{\text{TP}} \times (b \cdot s \cdot h \cdot 2)$$ - **结论**：$V_{\text{TP+SP}} = V_{\text{TP}}$。操作次数从 2 次变为 4 次，每次数据量减半；在机内 NVLink 上多出的两次启动延迟（约 $2\mu s$）很小。 - 作为回报，模型所有非 TP 区域的驻留激活显存被**直接消减了 $\text{TP}$ 倍**（例如 8 卡下直接砍掉 87.5% 激活显存）。 ## 具体演示与数值推演以一个 70B 模型（隐藏层维度 $h=8192$，层数 $L=80$），在 $\text{TP}=8$、序列长度 $s=16384$（16k 上下文）、单卡批次 $\text{MBS}=1$（BF16 精度，每元素 2 字节）下的单层驻留激活为例： - **未开启 SP（原生 TP-only）**：按 Korthikanti 等的激活公式，TP 下每层激活为 $sbh\left(10 + \frac{24}{t} + \frac{5as}{ht}\right)$ 字节，其中未被 TP 切分的部分约 $10sbh$：两个 LayerNorm 的输入（$4sbh$）、QKV 投影与 MLP 第一层的输入（$4sbh$）、两个 Dropout 掩码（$2sbh$，每元素 1 字节）。单层为： $$M_{\text{act/layer}} \approx 10 \times 16384 \times 1 \times 8192\text{ 字节} \approx 1.34\text{ GB}$$ 全网 80 层累加： $$M_{\text{total\_act}} \approx 80 \times 1.34\text{ GB} \approx 107.4\text{ GB}$$ 仅这部分未切分的激活就已超过单卡 80 GB，不开 SP 时只能依赖重计算。 - **开启 SP（TP + SP 协同）**：这些张量沿序列切为 $\frac{16384}{8} = 2048$ 个 token，单层降为： $$M_{\text{act/layer\_SP}} \approx 10 \times 2048 \times 1 \times 8192\text{ 字节} \approx 167.8\text{ MB} \quad (\text{为原来的 } 1/8)$$ 80 层合计约 $13.4\text{ GB}$，省出约 $94\text{ GB}$。 ## 代码填空题模拟一个 transformer block 内激活形状 `(s, h)` 的变化，比较 TP-only 与 TP+SP 的峰值激活。 ```{python} #| course_role: exercise def trace_block( seq: int, hidden: int, tp: int, enable_sp: bool, batch: int = 1, ) -> list[tuple[str, int, int]]: """ 返回 [(区域名, s, h)]：前向经过 block 时激活形状的轨迹（每 token 两字节 BF16 可算字节数）。算子表（对 (s, h) 的作用）： split_seq : s -> s/tp （SP 区入口 / g* 出口） gather_seq : s -> s*tp （g：SP -> TP） split_hidden : h -> h/tp （col-linear 切 h） gather_hidden: h -> h*tp （row-linear 恢复 h） noop : 不变区域顺序： layernorm → col_linear → gelu → row_linear → dropout """ def split_seq(s, h): return (s // tp, h) def gather_seq(s, h): return (s * tp, h) def split_hidden(s, h): return (s, h // tp) def gather_hidden(s, h): return (s, h * tp) def noop(s, h): return (s, h) region_ops = { "layernorm": ______, # 填空：TP-only 下 LN 输入就是全量 (s,h)，用什么算子？ "col_linear": ______, # 填空：col-linear 沿隐藏维切分 "gelu": ______, # 填空：逐元素，形状不变 "row_linear": ______, # 填空：row-linear 恢复隐藏维 "dropout": ______, # 填空：TP-only 下 dropout 保持全量 } sp_ops = { # 启用 SP 时各区域用的算子（提示：LN/dropout 在 SP 区沿序列切， # col/row 在 TP 区沿隐藏维切） "layernorm": ______, "col_linear": ______, "gelu": ______, "row_linear": ______, "dropout": ______, } plan = sp_ops if enable_sp else region_ops s, h = seq, hidden trace = [] for name, op in plan.items(): if enable_sp and name == "col_linear": # SP -> TP 过渡：g = all-gather 沿序列维，恢复完整 s s, h = gather_seq(s, h) trace.append(("g (all-gather seq)", s, h, False)) # 瞬态通信缓冲 s, h = op(s, h) # SP 模式下 row-linear 的 (s, h) 输出是瞬态（随后被 g* 切回序列分片）， # 反向真正需要的是它的输入（上一区域的 (s, h/tp)）——所以标记为不驻留。 persists = not (enable_sp and name == "row_linear") trace.append((name, s, h, persists)) # 驻留激活（要存给反向） if enable_sp and name == "row_linear": trace.append(("row_out (full h)", s, h, False)) # 瞬态：随后被 g* 切回 s, h = split_seq(s, h) trace.append(("g* (reduce-scatter seq)", s, h, False)) return trace def max_activation_bytes(trace: list[tuple[str, int, int, bool]], batch: int) -> int: """ 峰值【驻留】激活 = max(需要保存给反向的激活字节数)，BF16 每元素 2 字节。 g/g* 与 row_out 是瞬态全量张量（存在通信/临时缓冲区里），不计入驻留激活；但它们会出现在峰值显存里——面试时要主动区分"驻留激活"与"瞬态缓冲"。 """ return max( ______ # 填空：只统计驻留条目 (stored=True)，b*s*h*2 for name, s, h, stored in trace if stored ) def bf16_bytes(s: int, h: int, batch: int) -> int: return batch * s * h * 2 if __name__ == "__main__": seq, hidden, tp, bs = 4096, 4096, 8, 1 t_tp = trace_block(seq, hidden, tp, enable_sp=False, batch=bs) t_sp = trace_block(seq, hidden, tp, enable_sp=True, batch=bs) print("TP-only:") for name, s, h, stored in t_tp: tag = " " if stored else "*" print(f" {tag}{name:22s} s={s:5d} h={h:5d} {bf16_bytes(s, h, bs)/1e6:8.1f} MB") print("TP+SP（* = 瞬态全量，不计入驻留激活）:") for name, s, h, stored in t_sp: tag = " " if stored else "*" print(f" {tag}{name:26s} s={s:5d} h={h:5d} {bf16_bytes(s, h, bs)/1e6:8.1f} MB") m_tp = max_activation_bytes(t_tp, bs) m_sp = max_activation_bytes(t_sp, bs) print(f"\nTP-only 峰值 = {m_tp/1e6:.1f} MB = b·s·h·2 = {bs*seq*hidden*2/1e6:.1f} MB") print(f"TP+SP 峰值 = {m_sp/1e6:.1f} MB = b·s·h·2/tp = {bs*seq*hidden*2/tp/1e6:.1f} MB") assert m_sp * tp == m_tp print("验证：TP+SP 峰值激活 = TP-only 的 1/tp ✓") ``` ## 三个问答题 ### Q1 为什么 LayerNorm 不能像 MLP 那样沿隐藏维切分？请写出 LayerNorm 公式并指出哪一步依赖完整 h。dropout 为什么也需要 SP（提示：随机掩码的模式）？ ### Q2 教材用 f/f\* 描述 TP 区域的入口与出口（与非 TP 区域的边界）、用 g/g\* 描述 SP 与 TP 区域的交界。请分别写出前向和反向下 f、f\* 各是什么操作（no-op 还是 all-reduce），并解释为什么反向恰好互换。SP 区域为什么刻意不用 all-reduce？ ### Q3 教材说 TP+SP "通信操作数是 TP 的两倍（2 all-gather + 2 reduce-scatter vs 2 all-reduce），但总通信量等价"。请用第 3 课的结论（all-reduce = reduce-scatter + all-gather）证明这一点，并指出这为什么使 SP 几乎免费。 ## 检查与通过标准总分 10 分：代码正确 4 分（变换计划、g/g\* 过渡、峰值公式与验证）、三题各 2 分、通过线 8 分。一票否决项： - 把 SP 与 CP 混为一谈（本书定义下）。 - 认为 SP 沿隐藏维切分 LayerNorm 的输入。 - 说不出 g\* 是 reduce-scatter（把 row-linear 的求和与序列切分合并）或方向错误。 - 声称 SP 增加了通信量（而非通信次数增加、总量等价）。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。 ```{python} #| course_role: solution def trace_block( seq: int, hidden: int, tp: int, enable_sp: bool, batch: int = 1, ) -> list[tuple[str, int, int]]: """ 返回 [(区域名, s, h)]：前向经过 block 时激活形状的轨迹（每 token 两字节 BF16 可算字节数）。算子表（对 (s, h) 的作用）： split_seq : s -> s/tp （SP 区入口 / g* 出口） gather_seq : s -> s*tp （g：SP -> TP） split_hidden : h -> h/tp （col-linear 切 h） gather_hidden: h -> h*tp （row-linear 恢复 h） noop : 不变区域顺序： layernorm → col_linear → gelu → row_linear → dropout """ def split_seq(s, h): return (s // tp, h) def gather_seq(s, h): return (s * tp, h) def split_hidden(s, h): return (s, h // tp) def gather_hidden(s, h): return (s, h * tp) def noop(s, h): return (s, h) region_ops = { "layernorm": noop, # 填空：TP-only 下 LN 输入就是全量 (s,h)，用什么算子？ "col_linear": split_hidden, # 填空：col-linear 沿隐藏维切分 "gelu": noop, # 填空：逐元素，形状不变 "row_linear": gather_hidden, # 填空：row-linear 恢复隐藏维 "dropout": noop, # 填空：TP-only 下 dropout 保持全量 } sp_ops = { # 启用 SP 时各区域用的算子（提示：LN/dropout 在 SP 区沿序列切， # col/row 在 TP 区沿隐藏维切） "layernorm": split_seq, "col_linear": split_hidden, "gelu": noop, "row_linear": gather_hidden, "dropout": noop, } plan = sp_ops if enable_sp else region_ops s, h = seq, hidden trace = [] for name, op in plan.items(): if enable_sp and name == "col_linear": # SP -> TP 过渡：g = all-gather 沿序列维，恢复完整 s s, h = gather_seq(s, h) trace.append(("g (all-gather seq)", s, h, False)) # 瞬态通信缓冲 s, h = op(s, h) # SP 模式下 row-linear 的 (s, h) 输出是瞬态（随后被 g* 切回序列分片）， # 反向真正需要的是它的输入（上一区域的 (s, h/tp)）——所以标记为不驻留。 persists = not (enable_sp and name == "row_linear") trace.append((name, s, h, persists)) # 驻留激活（要存给反向） if enable_sp and name == "row_linear": trace.append(("row_out (full h)", s, h, False)) # 瞬态：随后被 g* 切回 s, h = split_seq(s, h) trace.append(("g* (reduce-scatter seq)", s, h, False)) return trace def max_activation_bytes(trace: list[tuple[str, int, int, bool]], batch: int) -> int: """ 峰值【驻留】激活 = max(需要保存给反向的激活字节数)，BF16 每元素 2 字节。 g/g* 与 row_out 是瞬态全量张量（存在通信/临时缓冲区里），不计入驻留激活；但它们会出现在峰值显存里——面试时要主动区分"驻留激活"与"瞬态缓冲"。 """ return max( max(batch * s * h * 2 for _, s, h, stored in trace if stored) # 填空：只统计驻留条目 (stored=True)，b*s*h*2 for name, s, h, stored in trace if stored ) def bf16_bytes(s: int, h: int, batch: int) -> int: return batch * s * h * 2 if __name__ == "__main__": seq, hidden, tp, bs = 4096, 4096, 8, 1 t_tp = trace_block(seq, hidden, tp, enable_sp=False, batch=bs) t_sp = trace_block(seq, hidden, tp, enable_sp=True, batch=bs) print("TP-only:") for name, s, h, stored in t_tp: tag = " " if stored else "*" print(f" {tag}{name:22s} s={s:5d} h={h:5d} {bf16_bytes(s, h, bs)/1e6:8.1f} MB") print("TP+SP（* = 瞬态全量，不计入驻留激活）:") for name, s, h, stored in t_sp: tag = " " if stored else "*" print(f" {tag}{name:26s} s={s:5d} h={h:5d} {bf16_bytes(s, h, bs)/1e6:8.1f} MB") m_tp = max_activation_bytes(t_tp, bs) m_sp = max_activation_bytes(t_sp, bs) print(f"\nTP-only 峰值 = {m_tp/1e6:.1f} MB = b·s·h·2 = {bs*seq*hidden*2/1e6:.1f} MB") print(f"TP+SP 峰值 = {m_sp/1e6:.1f} MB = b·s·h·2/tp = {bs*seq*hidden*2/tp/1e6:.1f} MB") assert m_sp * tp == m_tp print("验证：TP+SP 峰值激活 = TP-only 的 1/tp ✓") ``` ## 补全后的代码（关键填空处） ```python region_ops = { "layernorm": noop, # TP-only：LN 输入本来就是全量 (s, h) "col_linear": split_hidden, # col-linear 切 h "gelu": noop, # 逐元素 "row_linear": gather_hidden, # row-linear 恢复 h "dropout": noop, # TP-only：全量 } sp_ops = { "layernorm": split_seq, # SP 区：沿序列切 "col_linear": split_hidden, # TP 区：沿隐藏维切 "gelu": noop, "row_linear": gather_hidden, "dropout": noop, # SP 区：保持序列切分 } # 过渡：col_linear 前 g=gather_seq（SP→TP）；row_linear 后 g*=split_seq（TP→SP） # 驻留标记：g/g*/row_out 为瞬态（存在通信/临时缓冲区），不计入驻留激活 def max_activation_bytes(trace, batch): return max(batch * s * h * 2 for _, s, h, stored in trace if stored) ``` 运行结果（seq=4096, h=4096, tp=8）： ```python TP-only：LN/row/dropout 处形状 (4096, 4096) → 驻留峰值 33.6 MB = b·s·h·2 TP+SP ：layernorm (512,4096)、col/gelu (4096,512)、dropout (512,4096) → 驻留峰值 4.2 MB = b·s·h·2/tp ✓（g、row_out、g* 是瞬态全量）验证：TP+SP 驻留峰值 = TP-only 的 1/tp ✓ ``` ## 三个问答题要点 ### Q1 - LayerNorm：μ 与 σ² 沿 h 维求均值/方差（`LayerNorm(x) = γ·(x−μ)/√(σ²+ε)+β`），每 rank 只有 h/tp 无法算出正确的统计量 → 必须全量 h； - dropout：掩码是随机模式，各 rank 必须用相同种子生成一致的掩码（否则激活模式不同步），且掩码要覆盖全量张量； - 所以这些"便宜"操作反而是激活内存大户 → 沿序列维切分（SP）把它们也分片。 ### Q2 - f/f\*（TP 区域的入口/出口，不开 SP 时使用）：前向 f=no-op（输入已同步）、f\*=all-reduce（保证正确性）；反向互换（梯度形状是前向的转置，谁前向是 no-op 谁反向就是 all-reduce）； - g/g\*（SP 与 TP 交界）：前向 g=all-gather（序列拼回全量）、g\*=reduce-scatter（合并 row-linear 的部分和 + 切回序列分片）；反向为共轭互换； - SP 区域不用 all-reduce 的原因：all-reduce 需要物化全量激活（先 gather 再 reduce），峰值内存回到 b·s·h，SP 的初衷就没了；用 reduce-scatter 把"求和"与"切分"一步完成。 ### Q3 - 由 all-reduce = reduce-scatter + all-gather，纯 TP 的 2 次 all-reduce 与 TP+SP 的 2 次 all-gather + 2 次 reduce-scatter 字节数相同，只是操作数翻倍（推导详见第 4 节）。 - 额外代价只有实现复杂度和更多次启动延迟，带宽成本与 TP 相同，所以说 SP 几乎免费。 ## 通过要点 - 术语：本书 SP 专指 TP 的配套技术（LN/dropout 沿序列切）；独立的序列切分叫 CP（第 8 课）。 - 激活形状：TP 区域 h 切、SP 区域 s 切；**驻留**激活峰值 b·s·h/tp；g/g* 的瞬态全量张量只存在于通信/临时缓冲区（面试要主动区分）。 ## 官方主参考 - [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - [PyTorch distributed documentation](https://pytorch.org/docs/stable/distributed.html) ::: ::: {.course-pagination} [**← 上一课**<br>第 6 课：张量并行 TP](lesson06.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 8 课：上下文并行 CP](lesson08.qmd){.pagination-link .pagination-next} :::