训练长上下文时,参数量可能没变,但中间激活急剧变大。 对 GPT 类模型,常见 hidden activation 形状是 [S, B, H]。 当序列长度 S 从 4K 变成 128K,这部分显存线性放大 32 倍。 Attention 更麻烦:朴素 attention score 是 [B, heads, S, S], 这就是平方项。FlashAttention 能避免物化完整 score 矩阵,但 K/V、mask、dropout、反向所需状态仍然会让长序列成为系统问题。

一个显存估算脚本

def gib(x):
    return x / 2**30

def hidden_activation_gib(seq, batch, hidden, bytes_per_elem=2, copies=6):
    # copies 粗略表示 layernorm/dropout/residual/qkv 等多份中间状态
    return gib(seq * batch * hidden * bytes_per_elem * copies)

for seq in [4096, 16384, 65536, 131072]:
    mem = hidden_activation_gib(seq, batch=1, hidden=8192)
    print(f"S={seq:6d} hidden activation rough={mem:6.2f} GiB / layer group")

Megatron 风格的 Sequence Parallel 通常和 Tensor Parallel 一起使用。 Day19 看到 TP 会让 MLP/Attention 的某些中间张量按 hidden/head 维度切分, 但 LayerNorm、Dropout、Residual 这类操作常常仍然在每个 TP rank 上保存完整 [S, B, H] activation。 SP 的做法是:在这些不需要完整 hidden 分片通信的区域,把 activation 沿序列维度 S 切开, 每个 TP rank 只保存 S / TP 的 tokens。

SP 的收益和限制

# Megatron-LM 常见配置片段
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  --sequence-parallel \
  --seq-length 8192 \
  --micro-batch-size 1

# 直觉:
# TP=4 时,SP 区域里每个 rank 只常驻约 S/4 的 token activation。
# 需要完整序列参与 TP compute 的边界,再用 all-gather 取回。

Context Parallel 更直接地面对长上下文: 把完整序列 S 切成多个 context shard,每个 rank 只保存一段 tokens 的 Q/K/V 和中间状态。 但 self-attention 的语义要求每个 query 能看到全局 key/value。 所以 CP 的核心不是简单切数据,而是设计一种方式,让每个 rank 用本地 Q 和其它 rank 的 K/V 分块完成全局 attention,同时不把所有 K/V 永久 all-gather 到本地。

SP vs CP 快速对照

Ring Attention 的核心直觉很漂亮: 每个设备持有一段 query 和一段 key/value。 query 留在本地,key/value block 沿着设备环移动。 每来一个远端 K/V block,本地就做一次 blockwise attention 更新输出和 softmax 统计量。 在下一块 K/V 传输时,当前块的 attention 计算同时进行。 如果 overlap 做得好,通信可以被计算藏起来。

Ring Attention 的四步循环

为什么需要 blockwise softmax

attention 不能简单把每个 K/V block 的 softmax 输出相加。 softmax 的分母是全局所有 key 的指数和。 因此 blockwise attention 需要在线维护每一行的最大值和归一化分母, 类似 FlashAttention 的 online softmax。 这就是 Ring Attention 能保持精确语义的关键。

# 心智伪代码:不是可直接运行实现
for step in range(context_parallel_size):
    scores = q_local @ k_block.T
    scores = apply_causal_mask_for_this_block(scores, step)
    out_accum, m_i, l_i = online_softmax_update(
        out_accum, m_i, l_i, scores, v_block
    )
    send_async(k_block, v_block, next_rank)
    k_block, v_block = recv_async(prev_rank)

out_local = out_accum / l_i

今天的练习不要求真的训练百万 token,而是能从配置推导每张卡持有什么张量、 哪些边界需要 all-gather / reduce-scatter / ring send-recv, 以及长序列扩展后显存账怎么变化。

def sp_activation_per_rank(seq, batch, hidden, tp, bytes_per_elem=2):
    return seq // tp * batch * hidden * bytes_per_elem / 2**30

def cp_tokens_per_rank(seq, cp):
    return seq // cp

seq = 131072
hidden = 8192
batch = 1

for tp in [1, 2, 4, 8]:
    print(f"SP TP={tp}: hidden shard ~= {sp_activation_per_rank(seq, batch, hidden, tp):.2f} GiB")

for cp in [1, 2, 4, 8]:
    print(f"CP={cp}: tokens per rank = {cp_tokens_per_rank(seq, cp)}")

配置推导模板

1. 给定 S=131072,B=1,H=8192,BF16,估算一个 [S,B,H] activation 的大小,并说明为什么真实训练会有多份 copies。
2. 画出 Megatron Sequence Parallel 的形状流转:SP 区域 [S/TP,B,H]、TP compute 边界 all-gather、输出 reduce-scatter。
3. 用一张表区分 SP 和 CP:目标、切分对象、通信模式、适用场景。
4. 画出 4 rank Ring Attention:每个 rank 本地 Q 不动,K/V block 沿环移动,每一轮做什么计算和通信。
5. 解释为什么 blockwise attention 不能简单把每个 K/V block 的 softmax 输出相加,必须维护 online softmax 的全局统计。