AI大模型分布式训练：张量并行与流水线并行深度实践

引言

从GPT-3的175B参数到GPT-4的万亿参数级混合专家架构，大语言模型的参数规模每年以数量级速度增长。单张NVIDIA H100 80GB显存仅能存放约175B FP16参数（不含优化器状态），训练一个700B模型所需的硬件资源远超单台服务器。模型并行（Model Parallelism）——将神经网络不同层或同一层的不同张量切分到多块GPU上执行——已成为大规模训练不可或缺的三驾马车之一。

本文将深入拆解模型并行的两种核心策略：张量并行（Tensor Parallelism, TP）与Pipeline Parallelism（PP），从原理、实现、性能优化到实战落地的全链路进行解析。我们将以PyTorch + Megatron-LM为背景，穿插NCCL通信原语、GPU拓扑感知、微批次调度等工程细节，目标是让你在阅读后能在自有集群上跑出高效的分布式训练任务。

1. 三维并行体系：TP × PP × DP

在深入TP和PP之前，先建立全局视角。现代大模型训练通常采用三维并行策略：

• 数据并行（Data Parallelism, DP）：每个GPU持有完整模型副本，分片处理不同数据，梯度通过AllReduce同步。计算效率最高，但受限于单GPU内存。
• 张量并行（Tensor Parallelism, TP）：将单层内的权重矩阵按行或列切分到多GPU，每层前向/反向中通过AllGather/ReduceScatter通信。需要NVLink或PCIe通信。
• 流水线并行（Pipeline Parallelism, PP）：将模型的连续层按阶段（Stage）分配给各GPU，每个Stage处理不同Micro-batch，通过点对点通信传递激活值。

以NVIDIA Megatron-LM训练GPT-3 175B为例：TP=8（单节点内NVLink互联8卡），PP=16（16个Stage跨节点），DP=384×8/8/16=24路。总GPU数 = TP × PP × DP = 8×16×24 = 3072张V100。

2. 张量并行：矩阵分块的数学本质

2.1 核心原理

Transformer的核心计算是线性变换 Y = XW，其中 X 形状为 [batch, seq, hidden]，W 形状为 [hidden, ffn_hidden]（或反之）。张量并行的本质是将 W 按列切分：

# 列并行（Column Parallel）
# W [hidden, ffn] 被切成 k 份: W_0 [hidden, ffn/k], ..., W_{k-1}
# 各GPU计算 Y_i = X @ W_i，得到 Y 的第 i 片段
# 拼接后 Y = [Y_0 | Y_1 | ... | Y_{k-1}]

import torch
import torch.distributed as dist

class ColumnParallelLinear(torch.nn.Module):
    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        self.world_size = world_size
        # 只分配 out_features // world_size 列
        self.weight = torch.nn.Parameter(
            torch.randn(in_features, out_features // world_size)
        )
        self.bias = torch.nn.Parameter(
            torch.zeros(out_features // world_size)
        )
    
    def forward(self, x):
        # x: [batch, seq, in_features]，所有GPU上相同
        return torch.matmul(x, self.weight) + self.bias  # 无通信，完全并行

2.2 Row Parallel 与 AllReduce

注意：单纯列并行会在下一层引发问题。如果下一层是 Row Parallel（按行切分，每个GPU只算部分输出后需要求和），那么上一层不需要 Collect 输出；但如果后续需要完整的 Y 做残差连接（Y + X），则需要 AllGather。在 Megatron 的设计中，FFN 部分的列并行后紧跟行并行，用 AllReduce 替代 AllGather：

class RowParallelLinear(torch.nn.Module):
    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        self.weight = torch.nn.Parameter(
            torch.randn(in_features // world_size, out_features)
        )
    
    def forward(self, x_parallel):
        # x_parallel: 已经是列并行的输出，每个GPU持有不同切片
        partial = torch.matmul(x_parallel, self.weight)  # partial sum
        # AllReduce 求和得到最终 Y
        dist.all_reduce(partial, op=dist.ReduceOp.SUM)
        return partial  # [batch, seq, out_features]，所有GPU上相同

2.3 Attention 的 TP 拆分

Multi-Head Attention 天然适合 TP：12 头 Attention 在 4 路 TP 下，每个 GPU 仅计算 3 个 Attention 头。这意味着 Q/K/V 权重按列切，输出投影按行切，完全数学等价，不引入额外通信：

class TensorParallelAttention(torch.nn.Module):
    def __init__(self, hidden_size, num_heads, tp_size):
        assert num_heads % tp_size == 0
        self.num_heads_per_partition = num_heads // tp_size
        self.qkv = ColumnParallelLinear(hidden_size, hidden_size * 3, tp_size)
        self.output = RowParallelLinear(hidden_size, hidden_size, tp_size)
    
    def forward(self, x, mask):
        # QKV: 每个GPU拿到 [batch, seq, num_heads_per_partition * head_dim * 3]
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        # 各GPU独立计算自己负责的那些头 → Scaled Dot-Product Attention
        attn_out = scaled_dot_product_attention(q, k, v, mask)
        # AllReduce 聚合
        return self.output(attn_out)

2.4 通信开销分析

在 Transformer Block（Attention + FFN）中，TP 引入 4 次 AllReduce（Attention 输出 + FFN 输出，各含前向和反向）。每次 AllReduce 通信量为 2N/P（Ring-AllReduce 的两阶段传输），其中 N 是张量元素数，P 是 TP 并行度。

以 H100 NVLink 双向 900GB/s 带宽为例，70B 模型的 FFN 中间维度 28672 的激活值切分到 8 卡，单次 AllReduce 耗时约 1.04ms。这意味着在 TP=8 场景下，每层 Block 约引入 4ms 通信开销。若 GPU 计算延迟为 2ms（小的 Micro-batch），通信占比将高达 67% —— 这正是 TP 不能跨节点扩展的根本原因。

3. 流水线并行：时间换空间的流水线艺术

3.1 朴素 Pipeline 与 Bubble 问题

将模型按层均分到 k 个 Stage（每 Stage 驻留于不同 GPU），前向传播时依次执行 Stage 0 → Stage 1 → ... → Stage k-1。但这会产生严重的 Pipeline Bubble —— Stage i 在处理 Micro-batch m+1 之前必须等待 Stage i+1 对 Micro-batch m 的计算完成（依赖梯度的反向传播）。

朴素 GPipe 方案的 Bubble 率为 (k-1)/(m + k - 1)，其中 m 是 Micro-batch 数。当 k=16，m=48 时，Bubble 率约 23% —— 即 23% 的时间 GPU 在空转等待。

3.2 1F1B 调度：消除 Bubble 的利器

Megatron-LM 采用 1F1B（One Forward One Backward） 交错调度策略。核心思路：在完成一定数量的 Warmup 前向之后，每个 Stage 严格交错执行 1 次 Forward + 1 次 Backward，这样每个 Stage 最多同时缓存 m 个 Micro-batch 的激活值。

# 1F1B 调度伪代码 (每个Pipeline Stage独立执行)
def schedule_1f1b(stage_id, num_microbatches, num_stages):
    warmup_forward = num_stages - stage_id - 1
    
    # Phase 1: Warmup —— 只需前向
    for m in range(warmup_forward):
        forward(m)
    
    # Phase 2: 稳定态 —— 1 Forward + 1 Backward 交错
    for m in range(warmup_forward, num_microbatches):
        forward(m)
        backward(m - warmup_forward)
    
    # Phase 3: Cooldown —— 仅后向
    for m in range(num_microbatches - warmup_forward, num_microbatches):
        backward(m)

1F1B 将 Bubble 率从 (k-1)/(m+k-1) 降低到约 (k-1)/m。在 k=16，m=48 的场景下，Bubble 率降至约 3%。代价是每个 Stage 最多缓存 m 个 Micro-batch 的激活值（用于反向计算梯度），这恰好通过Activation Checkpointing控制在可接受范围。

3.3 Interleaved 1F1B（虚拟 Stage）

Chew 等人提出的 Interleaved 1F1B 进一步优化 Bubble：每个 GPU 不再是连续持有 "第 1-10 层"，而是交错持有多个 Virtual Stage。一个 GPU 持有层 [1, 9, 17, 25]（同时是 Stage 0 和 Stage 1 的一部分），这样的 Virtual Pipeline Stage 每个 GPU 参与 k 次前向（而非1次），使通信更均匀分布于时间轴上。

Bubble 率进一步降低为 (k-1)/(m × v)，其中 v 是 Virtual Stage 数量。代价是 GPU 间通信次数增加 v 倍，但对高速 NVLink 集群而言通常可接受。

3.4 Activation Checkpointing 与内存权衡

PP 的核心优势是单 GPU 内存仅 O(N/k)，但代价是需要保存中间激活用于反向。选择性 Activation Checkpointing 在 Attention 和 FFN Block 出口处只保存小部分激活（输入），反向时重新计算整个 Block：

from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def __init__(self, dim, heads, tp_size):
        self.attn = TensorParallelAttention(dim, heads, tp_size)
        self.mlp = TensorParallelMLP(dim, tp_size)
    
    def forward(self, x, mask):
        # Checkpoint 只保存 Block 输入 x，反向时重跑整个计算
        x = x + checkpoint(self.attn, x, mask, use_reentrant=False)
        x = x + checkpoint(self.mlp, x, use_reentrant=False)
        return x

此技术将内存从 O(L × B×S×D)（所有层激活）降至 O((L/k) × v × B×S×D)，以约 30% 的计算冗余换取近 k 倍的内存压缩。在实际生产训练中几乎必用。

4. TP × PP 联合策略：Megatron 的工程实现

4.1 拓扑感知的分配策略

在实际部署中，TP 和 PP 不是自由组合的，受 GPU 拓扑严格约束。以 NVIDIA DGX H100 节点（8 GPU，NVLink 900GB/s 全互联）为例：

• TP 置于节点内：8 卡 NVLink 全互联提供充足带宽（总量 7.2 TB/s），适合 TP=8。
• PP 跨节点：节点间通过 InfiniBand HDR 200Gbps（约 25 GB/s）互联，PP 的点对点通信量与同行数据量（激活值大小）成正比，远小于 TP 的 AllReduce 全局通信。

Megatron 的并行分组代码：

# 利用 torch.distributed 构建并行组
def initialize_model_parallel(
    tensor_model_parallel_size=1,
    pipeline_model_parallel_size=1,
):
    world_size = dist.get_world_size()
    
    # 构建 Pipeline 组（在循环中）
    for i in range(world_size // pipeline_model_parallel_size):
        ranks = list(range(
            i * pipeline_model_parallel_size,
            (i + 1) * pipeline_model_parallel_size
        ))
        group = dist.new_group(ranks)
        if dist.get_rank() in ranks:
            pipeline_group = group
    
    # 构建 TP 组（在 Pipeline 组内）
    for i in range(pipeline_model_parallel_size):
        ranks = list(range(i, world_size, pipeline_model_parallel_size))
        group = dist.new_group(ranks)
        if dist.get_rank() in ranks:
            tp_group = group

    return tp_group, pipeline_group

4.2 序列并行（Sequence Parallelism）：Megatron-LM 的秘技

当 TP 将隐藏维度切分到 k 张卡后，原本不属于 Dropout/LayerNorm 计算的维度（batch×seq）并没有被并行化。Megatron-LM 的 Sequence Parallelism 将 LayerNorm 和 Dropout 沿序列维度切分，使这些小型计算也获得 k 路并行，同时减少激活值内存（序列维度上的激活只需保留 1/k）。

4.3 分布式优化器（ZeRO Stage-1）

ZeRO-1 与 PP 天然配合：各 PP Stage 仅保存本 Stage 参数的优化器状态片段，通过 ReduceScatter 在反向结束时分发。这使得 70B 模型在 TP=8 PP=16 配置下，单 GPU 优化器显存占用从约 1.1TB（全量 Adam 状态）压缩到约 86GB，可放入单张 H100 80GB。

5. 性能调优：让集群跑出线性加速比

5.1 Micro-batch Size 的选择

Micro-batch size 是 PP 调优的核心杠杆。太小会导致计算碎片化，GPU利用率低；太大会导致激活缓存超出显存，且 Bubble 绝对时间增加。经验公式：m ≥ 4k（PP Stage 数的 4 倍以上）。

对于 Llama-3 70B，典型配置：TP=8（单节点），PP=8（8节点），Micro-batch=4，Global Batch=2048 tokens/step。

5.2 NCCL 通信优化

• NCCL_BUFFSIZE：默认 4MB，TP AllReduce 建议设为 32MB–64MB，可提升小消息场景吞吐。
• NCCL_PROTO=LL128：使用 LL128 协议替代默认的 Simple 协议，在大规模集群减少延迟。
• P2P 层次化通信：PP 点对点使用 dist.isend/dist.irecv 时，指定 group=pipeline_group 并使用 NVLink P2P 路径（通过 NCCL_P2P_LEVEL=NVL）。

5.3 Flash Attention 与 TP 的协同

FlashAttention 本身不改变 TP 的通信模式，但显著降低 Attention 计算时间，使通信隐藏更容易实现。在 TP=8 的场景中，FlashAttention 将注意力计算从 O(N²) HBM 读写优化到 O(N) SRAM 计算，单头计算延迟从 80ms 降至 12ms（seq_len=4096, head_dim=128, A100），使 AllReduce 更容易与下一层计算重叠。

6. 实战：PyTorch FSDP2 + TP 训练管線

6.1 使用 FSDP2 的 TP 并行化

PyTorch FSDP v2（基于 DeviceMesh）已原生支持 TP。对于 7B–14B 模型，FSDP2 + TP 是快速上手的好选择：

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, parallelize_module
)

# 初始化设备网格
mesh = init_device_mesh("cuda", (tp_size, pp_size), mesh_dim_names=("tp", "pp"))
tp_mesh = mesh["tp"]

# 对一个 Transformer Block 做 TP
parallelize_module(transformer_block, tp_mesh, {
    "self_attn.q_proj": ColwiseParallel(),
    "self_attn.k_proj": ColwiseParallel(),
    "self_attn.v_proj": ColwiseParallel(),
    "self_attn.o_proj": RowwiseParallel(),
    "mlp.gate_proj": ColwiseParallel(),
    "mlp.up_proj": ColwiseParallel(),
    "mlp.down_proj": RowwiseParallel(),
})

6.2 完整训练循环

# 完整的 TP+PP 训练循环（简化示意）
dist.init_process_group(backend="nccl")
mesh = init_device_mesh("cuda", (tp_size, pp_size), 
                        mesh_dim_names=("tp", "pp"))

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-70B")
# TP 并行化每一层
for layer in model.model.layers:
    parallelize_module(layer, mesh["tp"], TP_STYLE)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = torch.cuda.amp.GradScaler()  # BF16 混合精度

for step, batch in enumerate(dataloader):
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        outputs = model(**batch, labels=batch["input_ids"])
        loss = outputs.loss
    
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

6.3 调试 Checklist

• 梯度一致性：在 TP=PP=1 单卡上跑一组数据记录 Loss 曲线，对比多卡下 Loss 是否在 1e-6 误差内一致 —— 这是排查通信错误的黄金标准。
• 通信 Profiling：使用 nsys profile --trace=cuda,nvtx 抓取训练 Profile，重点观察 AllReduce 与 Compute 是否有效重叠。
• 内存 OOM：逐步降低 Micro-batch Size，同时检查 Activation Checkpointing 是否真正启用了 PP Stage 出口。
• NCCL 超时：大规模集群设置 NCCL_TIMEOUT=1800（秒），避免偶发抖动被默认超时踢出组。

7. 未来趋势：TP/PP 的演进与新兴并行维度

模型并行不仅是训练基石，也开始渗透推理部署。NVIDIA TensorRT-LLM 在推理引擎中引入 TP，将 70B 模型的 KV Cache 切分到 8 卡，突破单卡 80GB 显存限制，支持 100K+ 上下文长度。

展望未来几个关键方向：

• Context Parallelism (CP)：将超长序列切分到多 GPU（不同于 PP 按层切分），是 Llama-3 100K 上下文的关键技术（Ring Attention）。
• Expert Parallelism (EP)：MoE 模型中不同 Expert 分配到不同 GPU，DeepSeek-V3 的生产级方案将 256 个 Routed Expert 分布到 256 张 H100 上。
• 光互联 PP：Pentink 等研究用光交换替代电互联 PP P2P，带宽提升 10×，Bubble 率降至接近零。
• 异构内存池：结合 CXL 内存扩展与 PP，让 Stage 间共享远端内存，进一步降低单 GPU 显存需求。

结语

张量并行和流水线并行不是高深的学术黑箱，而是矩阵分块、通信原语与 GPU 拓扑约束的精确组合。Megatron-LM 的成功不在于提出了全新算法，而在于将 Ring-AllReduce、1F1B 调度、Activation Checkpointing 和 Sequence Parallelism 这些久经考验的工程实践整合为一个内聚系统。

下次面对 "模型太大，单卡放不下" 时，记住：TP 解决 FFN 的 GPU 内存问题，PP 解决 Transformer 层数的 GPU 内存问题，两者结合撑起万亿参数大模型的生产训练。 而这两种策略选择的本质，始终是在通信开销、计算效率与内存约束之间寻找最优平衡点。