TP / PP 的 PyTorch 原生实现

TorchTitan 不依赖任何外部分布式库（如 Megatron / ColossalAI），TP / PP / SP / CP 全靠 PyTorch 原生 API 实现。本章看每种的写法。

6.1TP（Tensor Parallel）

用 parallelize_module() + ParallelStyle。看 TorchTitan 的 LlamaAttention TP plan：

from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, parallelize_module
)

attn_plan = {
    "wq": ColwiseParallel(),
    "wk": ColwiseParallel(),
    "wv": ColwiseParallel(),
    "wo": RowwiseParallel(),
}
mlp_plan = {
    "w1": ColwiseParallel(),
    "w3": ColwiseParallel(),
    "w2": RowwiseParallel(),
}
parallelize_module(layer.attention, mesh["tp"], attn_plan)
parallelize_module(layer.feed_forward, mesh["tp"], mlp_plan)

跟 Megatron 思路完全一样：QKV 列切、O 行切；gate/up 列切、down 行切。区别是 PyTorch 内部直接用 DTensor 表示切完的 weight，不需要换 Linear 类。

6.2SP（Sequence Parallel）

SP 是 TP 的扩展：把 layernorm / dropout / residual 沿 seq 维切（这些是 TP 内部仍然 replicate 的部分）。

from torch.distributed.tensor.parallel import SequenceParallel

# 给 norm 加 SP
parallelize_module(layer, mesh["tp"], {
    "attention_norm": SequenceParallel(),
    "ffn_norm": SequenceParallel(),
    # attention/MLP plan 同上
})

TorchTitan 的 SP 跟 Megatron 的 split_gather SP 一样思路。

6.3PP（Pipeline Parallel）

用 torch.distributed.pipelining（2.5+ 稳定）。看简化：

from torch.distributed.pipelining import pipeline, ScheduleGPipe, Schedule1F1B

# 1) 切 model
def split_spec(model, mb):
    # 这里说明每个 stage 应该有哪些 layer
    return {
        "layers.0": "stage_0",
        "layers.8": "stage_1",
        ...
    }

stages = pipeline(model, mb_args=(input_ids,), split_spec=split_spec)

# 2) 选 schedule
schedule = ScheduleGPipe(stages, n_microbatches=8)
# 或 1F1B
schedule = Schedule1F1B(stages, n_microbatches=8)

# 3) 训
loss = schedule.step(input_ids, target=labels)

TorchTitan 自家代码在 torchtitan/distributed/pipeline.py，按 layer 数自动均分。

6.4CP（Context Parallel）—— 长上下文专用

CP 是 PyTorch 2.6+ 新加：把 attention 的 KV 沿 seq 维切（不切 head），用环状传 KV。等价于 Megatron 的 ring attention。

from torch.distributed.tensor.experimental import context_parallel

with context_parallel(mesh["cp"], buffers=[input_ids, labels]):
    out = model(input_ids)
    loss = criterion(out, labels)
loss.backward()

对应 toml：

[experimental]
context_parallel_degree = 4

典型用法：训 32K-128K 上下文时启用，CP × TP × FSDP 联合用。

6.5四种并行混搭：5D mesh

TorchTitan 默认 5 维 mesh：

(dp_replicate, dp_shard, cp, tp, pp)
   ↑节点间    ↑节点内   ↑长上下文  ↑权重切  ↑层切

举例：32 GPU 训 70B
- pp = 2
- tp = 4
- cp = 1
- dp_shard = 4
- dp_replicate = 1
- 32 = 1 × 4 × 1 × 4 × 2  ✓

TorchTitan 自动校验 nproc 是否等于各维乘积。

6.6各方案选择

6.7跟 Megatron 性能对比

Llama-3 70B、64×H100、bf16、seqlen=8K：

6.8这章你需要带走的

• TP 用 parallelize_module + Col/Row Parallel；
• SP 用 SequenceParallel ParallelStyle；
• PP 用 torch.distributed.pipelining；
• CP 用 context_parallel() context manager；
• 5D mesh：(dp_replicate, dp_shard, cp, tp, pp)。