Sequence、Context、Pipeline 与 Expert Parallel

先按“切什么”分类

这些轴可形成多维 mesh，例如 TP=2、CP=4、PP=2、EP=8；但 process groups、rank coordinates 与 collective 顺序必须明确。总 GPU 数不一定是简单乘积：某些 dimensions 只用于特定 layers，EP group 内还可嵌套 TP。

GPU 互联拓扑进一步把这些逻辑 groups 映射到 NVLink/NVSwitch/PCIe/NIC/NUMA 路径；同一并行度不等于相同通信成本。

Sequence Parallel：逐 token 算子天然可本地执行

对 X:[B,T,H] 沿 $T$ 切成 $p$ 份，RMSNorm 每个 token 只沿 $H$ 归约，所以：

这可降低每 rank 在 norm/residual/dropout 等区域的 activation memory。它不让 attention 自动只看本地 tokens：若下一算子需要完整 sequence 或采用不同 shard layout，必须通信转换。训练中 SP 还减少保存 activations；本页聚焦推理 ownership。

Context Parallel：softmax 必须跨 context shards 精确合并

设单 Query、K/V 沿 source length 分到 ranks：$K=[K^{(0)};K^{(1)}]$。每 rank 计算 local scores $s^{(r)}$，维护：

全局最大值：

重缩放后合并：

这与 FlashAttention online softmax 的 tile 合并同构，只是 partials 位于不同 ranks。具体 CP 可 all-gather K/V、ring 传 K/V blocks、或通信 partial stats/output；逻辑必须覆盖全局可见 context 并正确处理 causal position。

CP 减少每 rank KV payload到约 $1/p$，但每个 Query 仍需要全局 context 的等价贡献，通信/同步与拓扑成为成本。单 GPU 无法验证 NCCL 性能。

RoPE、PI、NTK-aware 与 YaRN定义每个全局 logical position 的旋转频率；CP 只切分这些 positions/KV 的 rank 所有权。各 rank 若使用不同 RoPE schedule 或把 local offset 误当 global position，softmax partial 合并即使公式正确也会得到错误结果。

Pipeline Parallel：切层并引入 bubble

若 $S$ 个 stages、$M$ 个 microbatches，仅考虑 forward、每 stage/microbatch 同成本，简单 pipeline ticks 为：

stage-slot 利用率：

$S=3,M=4$ 时 schedule：

tick 0: [mb0, idle, idle]
tick 1: [mb1, mb0, idle]
tick 2: [mb2, mb1, mb0]
tick 3: [mb3, mb2, mb1]
tick 4: [idle, mb3, mb2]
tick 5: [idle, idle, mb3]

利用率为 $4/(4+3-1)=2/3$。实际推理 stage 耗时不均、prefill/decode shape 不同、microbatch 到达动态、activation transfer 和 KV ownership 都会改变 bubble；不能用该公式承诺线性扩展。Continuous batching/P-D disaggregation 下还有更复杂 schedule。

可运行的三个不变量

import math,platform,torch

torch.manual_seed(223)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float64

def rms_norm(x,weight,eps=1e-6):
    return x*torch.rsqrt(x.square().mean(-1,keepdim=True)+eps)*weight

def merge_context_partials(q,k_shards,v_shards):
    local=[]
    for k,v in zip(k_shards,v_shards):
        scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
        m=scores.max(-1).values; exp=torch.exp(scores-m[...,None])
        local.append((m,exp.sum(-1),exp@v))
    global_m=torch.stack([part[0] for part in local]).max(0).values
    denominator=sum(part[1]*torch.exp(part[0]-global_m) for part in local)
    numerator=sum(part[2]*torch.exp(part[0]-global_m)[...,None] for part in local)
    return numerator/denominator[...,None]

def pipeline_schedule(stages,microbatches):
    total=stages+microbatches-1; schedule=[]; active=0
    for tick in range(total):
        row=[]
        for stage in range(stages):
            mb=tick-stage
            if 0<=mb<microbatches: row.append(f"mb{mb}"); active+=1
            else: row.append("idle")
        schedule.append(row)
    return schedule,active/(stages*total)

def main():
    x=torch.randn(2,8,6,device=DEVICE,dtype=DTYPE)
    weight=torch.randn(6,device=DEVICE,dtype=DTYPE)
    dense_norm=rms_norm(x,weight); shards=x.chunk(2,dim=1)
    sequence_parallel=torch.cat([rms_norm(shard,weight) for shard in shards],dim=1)
    torch.testing.assert_close(sequence_parallel,dense_norm)

    q=torch.randn(1,2,1,4,device=DEVICE,dtype=DTYPE)
    k=torch.randn(1,2,8,4,device=DEVICE,dtype=DTYPE); v=torch.randn_like(k)
    dense=torch.softmax(q@k.transpose(-2,-1)/math.sqrt(q.shape[-1]),-1)@v
    context=merge_context_partials(q,k.chunk(2,dim=-2),v.chunk(2,dim=-2))
    torch.testing.assert_close(context,dense,rtol=1e-10,atol=1e-10)

    schedule,utilization=pipeline_schedule(3,4)
    assert schedule==[["mb0","idle","idle"],["mb1","mb0","idle"],
        ["mb2","mb1","mb0"],["mb3","mb2","mb1"],
        ["idle","mb3","mb2"],["idle","idle","mb3"]]
    assert abs(utilization-2/3)<1e-12
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE}")
    print(f"sequence_parallel input={tuple(x.shape)} shards={[tuple(s.shape) for s in shards]} max_diff={(sequence_parallel-dense_norm).abs().max():.3e}")
    print(f"context_parallel Q={tuple(q.shape)} K_shards={[tuple(s.shape) for s in k.chunk(2,dim=-2)]} max_diff={(context-dense).abs().max():.3e}")
    print(f"pipeline_schedule={schedule}")
    print(f"pipeline_stage_slot_utilization={utilization:.3f}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/parallelism_overview.py。实际 sequence-parallel RMSNorm 最大差 0，context partial merge 最大差 1.110e-16，pipeline schedule 利用率 0.667。

脚本在单进程中模拟 shards，不证明通信 backend、overlap 或多 GPU 性能。它证明三个数学/调度不变量；生产 CP 要用 NCCL/拓扑 trace，PP 要测 stage balance 与真实 microbatch arrival。

成本与组合失败场景

• SP：减少逐 token activation replication，但增加 layout transitions；若频繁 All-Gather 后马上再切分，通信/内存峰值抵消收益。
• CP：减少每 rank KV 容量，增加 context exchange/softmax merge；短上下文或慢互联可能不划算。
• PP：减少每 rank weights，增加 stage transfer 与 bubble；单请求 decode microbatch 少时利用率尤其受限。
• EP：只存本 rank experts，增加 route imbalance 与两次 All-to-All。

组合时最常见的工程错误是 process-group coordinate 混淆：同一 global rank 同时属于 TP/CP/PP/EP groups，collective 次序必须只在对应 group 内一致。权重、KV、activation、router metadata 要分别标注 replicated/sharded/partial；“rank 3 有这个 tensor”不足以说明是哪一维 shard。

模型加载必须按并行 mesh 放置 layer/expert/weight shards；Request Lifecycle取消请求时不能破坏跨 stages/ranks 的 collective 顺序；Benchmark 方法要求把本机语义实验与真实多卡拓扑性能分开。

参考资料

• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models — Sequence Parallelism。
• NVIDIA Megatron Core — Context Parallelism
• Huang et al., GPipe
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM