大模型训练并行技术全解析

📦 数据并行 Data Parallelism

核心思想：将训练数据切分成多个小批次，每张 GPU 上保存一份完整的模型副本，独立处理各自的数据。各卡独立计算梯度后，通过 All‑Reduce 通信同步梯度并取平均，最后所有卡用相同的平均梯度更新模型权重。

通信模式：All‑Reduce（梯度同步），通信量与模型参数量成正比。
优点：实现简单，是提高训练吞吐量的最直接手段。
缺点：每张卡都要存储完整模型、梯度和优化器状态，显存冗余严重。
适用场景：模型能单卡装下，且希望快速扩展数据吞吐量。
代表框架：PyTorch DDP, Horovod, DeepSpeed ZeRO (FSDP)。

graph LR subgraph "数据并行 (DP)" A[数据批次] --> B1[GPU 0
完整模型] A --> B2[GPU 1
完整模型] A --> B3[GPU N
完整模型] B1 --> C1[梯度1] B2 --> C2[梯度2] B3 --> C3[梯度N] C1 & C2 & C3 --> D[All-Reduce
平均梯度] D --> E1[更新权重] D --> E2[更新权重] D --> E3[更新权重] end

                    PyTorch DDP
                    import torch
                        import torch.distributed as dist
                        from torch.nn.parallel import DistributedDataParallel as DDP

                        dist.init_process_group("nccl")
                        model = MyModel().cuda()
                        model = DDP(model)   # 自动包装

                        for data, target in dataloader:
                            data, target = data.cuda(), target.cuda()
                            output = model(data)
                            loss = criterion(output, target)
                            loss.backward()          # 反向传播，DDP自动同步梯度
                            optimizer.step()
                            optimizer.zero_grad()
                    
                

🧠 张量并行 Tensor Parallelism

核心思想：将模型 同一层（如 Transformer 的注意力头或前馈网络）的权重矩阵按列或行切分到多张 GPU 上，各卡计算自己负责的部分，最后通过通信合并结果。最典型的实现是 Megatron‑LM 的 1D 张量并行。

通信模式：前向和反向均需要 All‑Reduce 或 All‑Gather，通信频繁且数据量大。
优点：能解决单层超大（如超大隐藏层）的显存瓶颈，精细分配计算。
缺点：通信开销极大，通常只能在单机内部（NVLink 高速互联）使用。
适用场景：单卡无法容纳单个大层（例如 10B+ 参数的 FFN 层）。
代表框架：Megatron‑LM, PyTorch 原生 TP（实验性）。

graph TD subgraph "张量并行 (TP) 示例：MLP 切列" X[输入 X] --> A1[GPU 0: W1 列切片] X --> A2[GPU 1: W1 列切片] A1 --> B1[线性部分1] A2 --> B2[线性部分2] B1 & B2 --> C[All-Reduce 合并] C --> D[激活函数] end

                    Megatron‑LM 风格 (概念)
                    # 假设 hidden_size = 4096, 切分为2卡 (每卡负责 2048 维)
                        class ColumnParallelLinear(nn.Module):
                            def __init__(self, in_features, out_features, world_size):
                                super().__init__()
                                self.out_features_per_part = out_features // world_size
                                self.weight = nn.Parameter(torch.empty(in_features, self.out_features_per_part))
                                # ... 前向: 本地线性变换，然后 All-Gather 或 All-Reduce 合并

                        # 在模型中使用：
                        self.linear1 = ColumnParallelLinear(hidden_size, hidden_size*4, world_size=2)
                        self.linear2 = RowParallelLinear(hidden_size*4, hidden_size, world_size=2)
                    
                

⛓️ 流水线并行 Pipeline Parallelism

核心思想：将模型按层分成多个阶段（Stage），每个 GPU 负责一个或多个连续层的计算。数据像流水线一样依次经过各阶段，前一个阶段的输出直接作为后一个阶段的输入。为了提升效率，通常使用 微批次（Micro‑batch） 流水线来减少空闲时间。

通信模式：阶段间 点对点 (P2P) 传递激活值和梯度，通信量相对较小。
优点：适合跨多机扩展（通信不密集），能解决模型过深的问题。
缺点：存在 流水线气泡（Bubble）（设备空闲），且需要处理微批次的调度。
适用场景：模型层数非常多（如 100+ 层），且希望跨多服务器部署。
代表框架：GPipe, PipeDream, PyTorch 的 `torch.distributed.pipeline.sync`。

sequenceDiagram participant GPU0 as GPU 0
层1-3 participant GPU1 as GPU 1
层4-6 participant GPU2 as GPU 2
层7-9 Note over GPU0,GPU2: 微批次 1,2,3 流水线执行 GPU0->>GPU1: 微批次1 前向 GPU0->>GPU1: 微批次2 前向 GPU1->>GPU2: 微批次1 前向 GPU0->>GPU1: 微批次3 前向 GPU1->>GPU2: 微批次2 前向 GPU2-->>GPU1: 微批次1 反向 GPU1-->>GPU0: 微批次1 反向 Note over GPU2: 存在少量气泡

                    PyTorch Pipe (概念)
                    from torch.distributed.pipeline.sync import Pipe

                        # 假设 model 是 nn.Sequential 分成多个 stage
                        model = nn.Sequential(layer1, layer2, layer3, ...)
                        model = Pipe(model, chunks=4)   # chunks = 微批次数量

                        for data, target in dataloader:
                            data = data.cuda()
                            output = model(data)        # 自动流水线并行
                            loss = criterion(output, target)
                            loss.backward()
                            optimizer.step()
                    
                

📜 序列并行 Sequence Parallelism

核心思想：将输入序列按 长度维度 切分到多张 GPU 上，以降低单卡处理长序列时的激活显存占用。早期 SP 是 TP 的补充（例如 Megatron‑LM 将 LayerNorm 和 Dropout 在序列维度切分），更现代的 SP（如 DeepSpeed Ulysses）则通过 All‑to‑All 通信在序列维度和注意力头维度间转换，优化长序列的注意力计算。

通信模式：All‑Gather / Reduce‑Scatter（Megatron 风格）或 All‑to‑All（Ulysses 风格）。
优点：显著降低激活显存，支持超长上下文窗口（如 128k tokens）。
缺点：引入了额外的通信开销，且需要配合 TP 或 DP 使用。
适用场景：处理超长序列（如文档级、长视频）的训练。
代表框架：Megatron‑LM (SP), DeepSpeed Ulysses, ColossalAI。

graph LR subgraph "序列并行 (SP) 示意" S[长序列
L tokens] --> S1[GPU 0: 前 L/2] S --> S2[GPU 1: 后 L/2] S1 --> A1[本地 Attention] S2 --> A2[本地 Attention] A1 & A2 --> C[All‑Gather 拼接
或 All‑to‑All] C --> O[输出完整序列] end

                    Megatron‑LM 序列并行 (概念)
                    # 在 Transformer 层中，将 LayerNorm 和 Dropout 沿序列维度切分
                        class ParallelTransformerLayer(nn.Module):
                            def __init__(self, ...):
                                self.input_layernorm = LayerNorm(hidden_size, sequence_parallel=True)
                                self.attention = Attention(...)   # 使用 TP 的注意力
                                self.post_attention_dropout = Dropout(..., sequence_parallel=True)
                                # 前向时，每个 GPU 只处理自己序列分片上的激活值
                                # 但注意力计算内部会通过 TP 通信合并头
                    
                

🧩 专家并行 Expert Parallelism

核心思想：专为 混合专家模型 (MoE) 设计。MoE 模型包含多个“专家”网络（通常是 FFN）和一个“门控”网络。EP 将不同的专家分配到不同的 GPU 上，每个 Token 通过门控网络路由到最合适的几个专家进行计算，最后聚合结果。

通信模式：Token 在专家间路由产生巨大的 All‑to‑All 通信开销。
优点：使得 MoE 模型可以在不显著增加计算量的情况下，大幅扩展参数量（如 1.6T 参数）。
缺点：通信复杂，负载均衡难以保证（需辅助损失函数），且对显存和带宽要求极高。
适用场景：超大规模 MoE 模型（如 Switch Transformer, GLaM, DeepSeek‑MoE）。
代表框架：DeepSpeed (MoE), Megatron‑LM (MoE), GShard。

                    DeepSpeed MoE (配置示例)
                    # 在 DeepSpeed 配置文件中启用 MoE
                        {
                            "train_batch_size": 1024,
                            "fp16": {"enabled": true},
                            "model_parallel": {
                                "moe": {
                                    "enabled": true,
                                    "ep_size": 4,            # 专家并行度 (GPU数)
                                    "num_experts": 16,
                                    "top_k": 2,
                                    "capacity_factor": 1.5,
                                    "moe_param_group": true
                                }
                            }
                        }

                        # 模型定义中使用 deepspeed.moe.layer.MoE 层
                    
                