本文由 AI 生成，可能有误

在当今大模型（Large Model）时代，模型规模的急剧增长对计算资源提出了前所未有的挑战。单张计算卡（如GPU）的算力和显存已远不能满足千亿甚至万亿参数模型训练与部署的需求。为了应对这一挑战，分布式训练与推理技术应运而生，通过将计算任务和模型参数分解到多张乃至成千上万张计算卡上，实现了对庞大模型的有效处理。

在分布式策略中，数据并行（Data Parallelism, DP）、专家并行（Expert Parallelism, EP）、张量并行（Tensor Parallelism, TP）、流水线并行（Pipeline Parallelism, PP）和序列并行（Sequence Parallelism, SP）是五种最为核心和常见的技术。理解它们的原理、优劣及适用场景，对于高效进行大模型研发与部署至关重要。

本文将以浅显易懂的方式，为您详细拆解这五种并行策略，并附上基于 PyTorch 的简化代码示例，助您深入理解并掌握这些关键技术。

1. 数据并行 (Data Parallelism, DP)：最常用的大规模训练策略

核心思想： 数据并行是最直观、最常用的并行方式。其核心思想是“模型复制，数据切分”。简单来说，就是将整个模型完整地复制到每一个计算设备（GPU）上，然后将训练数据（一个大的 batch）切分成多个小的 mini-batch，每个设备分配一个 mini-batch 进行独立计算。

工作流程：

1. 分发 (Scatter)： 在每个训练迭代开始时，主进程（通常是 rank 0）会将模型参数广播（broadcast）到所有参与训练的 GPU 上，确保每个 GPU 上的模型初始状态一致。同时，将一个大的数据 batch 切分成多个小的 mini-batch。
2. 前向传播 (Forward Pass)： 每个 GPU 使用自己分配到的 mini-batch 数据，独立地进行前向传播，计算出各自的损失（loss）。
3. 反向传播 (Backward Pass)： 每个 GPU 根据自己计算的损失，独立地进行反向传播，计算出模型参数的梯度（gradients）。
4. 梯度聚合 (Gradient Aggregation)： 这是数据并行的关键步骤。所有 GPU 需要将各自计算出的梯度进行聚合。最常用的聚合操作是 All-Reduce。All-Reduce 操作会收集所有 GPU 上的梯度，将它们相加（或求平均），然后再将聚合后的结果分发回每个 GPU。这样，每个 GPU 都拥有了基于整个大 batch 数据计算出的全局梯度。
5. 参数更新 (Optimizer Step)： 每个 GPU 使用聚合后的全局梯度，独立地更新自己的模型参数。由于所有 GPU 使用相同的全局梯度进行更新，因此它们的模型参数在每个迭代结束后仍然保持一致。

图解DP：

graph TD
    subgraph "训练迭代开始"
        A[模型参数] -->|复制到所有GPU| B1[GPU 1: 完整模型]
        A -->|复制到所有GPU| B2[GPU 2: 完整模型]
        A -->|复制到所有GPU| B3[GPU ...: 完整模型]
        D[整个Batch数据] -->|切分| D1[Mini-batch 1]
        D -->|切分| D2[Mini-batch 2]
        D -->|切分| D3[Mini-batch ...]
    end
    subgraph "并行计算"
        D1 --> B1
        D2 --> B2
        D3 --> B3
        B1 -->|计算梯度1| G1
        B2 -->|计算梯度2| G2
        B3 -->|计算梯度...| G3
    end
    subgraph "梯度聚合与更新"
        G1 & G2 & G3 -->|All-Reduce操作| AG[聚合后的全局梯度]
        AG -->|分发回所有GPU| B1
        AG -->|分发回所有GPU| B2
        AG -->|分发回所有GPU| B3
        B1 -->|更新模型参数| U1[更新后的模型1]
        B2 -->|更新模型参数| U2[更新后的模型2]
        B3 -->|更新模型参数| U3[更新后的模型...]
    end

优点：

• 实现简单： 主流深度学习框架（如 PyTorch 的 DistributedDataParallel (DDP) 和 TensorFlow 的 MirroredStrategy）都提供了高度封装的 API，用户只需几行代码即可实现数据并行。
• 通信开销相对固定： 主要的通信开销在于 All-Reduce 操作，其通信量与模型大小成正比，与 GPU 数量关系不大（在高效实现下），因此在一定范围内扩展性良好。
• 负载均衡： 只要每个 mini-batch 的大小和计算量相近，各个 GPU 的负载就是均衡的。

缺点：

• 显存冗余： 每个 GPU 都需要存储一份完整的模型参数、梯度和优化器状态，对于参数规模巨大的模型，这会极大地消耗显存。一个拥有 1750 亿参数的模型，其 FP32 参数本身就需要 700GB 显存，这远远超出了单张 GPU 的承载能力。
• 无法解决单模型过大的问题： 当模型大到单张 GPU 无法容纳时，数据并行便无能为力。

简单的 PyTorch DDP 代码示例：

以下是一个极简的 PyTorch DistributedDataParallel (DDP) 示例。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os

def setup(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    """清理分布式环境"""
    dist.destroy_process_group()

def train(rank, world_size):
    # 1. 初始化分布式环境
    setup(rank, world_size)

    # 2. 定义一个简单的模型
    model = nn.Linear(10, 5).to(rank)
    # 关键步骤：使用DDP包装模型
    ddp_model = DDP(model, device_ids=[rank])

    # 3. 定义损失函数和优化器
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # 4. 模拟训练过程
    for _ in range(10):
        # 模拟输入数据，每个GPU上的数据不同
        inputs = torch.randn(20, 10).to(rank)
        labels = torch.randn(20, 5).to(rank)

        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward() # DDP会自动进行梯度聚合
        optimizer.step()

        if rank == 0:
            print(f"Rank {rank}, Loss: {loss.item()}")

    # 5. 清理
    cleanup()

if __name__ == "__main__":
    # 假设我们有2个GPU
    world_size = 2
    # 在实际应用中，会使用 torch.multiprocessing.spawn 来启动多进程
    # 这里为了简化，只展示核心逻辑
    # train(rank=0, world_size=2)
    # train(rank=1, world_size=2)
    # 实际启动命令通常是：torchrun --nproc_per_node=2 your_script.py
    # 以下为概念性演示
    print("这是一个DDP的简化示例，实际运行需要通过torchrun或类似工具启动。")

2. 张量并行 (Tensor Parallelism, TP)：切分模型内部的“大”张量

核心思想： 当模型的单个层（特别是巨大的线性层或注意力头）的权重矩阵大到单张 GPU 无法容纳时，张量并行就派上了用场。它的核心思想是“层内并行”，即将一个大的张量（如权重矩阵）切分到多个 GPU 上，每个 GPU 只负责计算张量的一部分。

工作流程（以 Transformer 中的 MLP 层为例）：

一个标准的 Transformer MLP 块通常包含两个线性层，例如 Y = GeLU(X @ A) @ B，其中 X 是输入，A 和 B 是权重矩阵。

1. 列并行 (Column Parallelism) - 切分第一个线性层：

• 将权重矩阵 A 按列切分，例如切成 [A1, A2]，分别放到 GPU 1 和 GPU 2 上。
• 输入 X 被复制到两个 GPU 上。
• GPU 1 计算 Y1 = X @ A1，GPU 2 计算 Y2 = X @ A2。
• 此时，每个 GPU 只拥有部分结果 [Y1, Y2]。

• 将权重矩阵 B 按行切分，例如切成 [B1; B2]（分号表示行拼接），分别放到 GPU 1 和 GPU 2 上。
• GPU 1 用 Y1 计算 Z1 = Y1 @ B1，GPU 2 用 Y2 计算 Z2 = Y2 @ B2。
• 为了得到最终的结果 Z = Z1 + Z2，需要进行一次 All-Reduce 操作，将 Z1 和 Z2 相加并同步到所有 GPU。

图解TP：

graph TD
    subgraph "列并行 (Column Parallelism)"
        X[输入 X] --> GPU1 & GPU2
        A[权重 A] -->|按列切分| A1 & A2
        A1 --> GPU1
        A2 --> GPU2
        GPU1 -- 计算 X @ A1 --> Y1
        GPU2 -- 计算 X @ A2 --> Y2
        Y1 & Y2 -->|GeLU激活函数(逐元素)| GY1 & GY2
    end
    subgraph "行并行 (Row Parallelism)"
        B[权重 B] -->|按行切分| B1 & B2
        B1 --> GPU1
        B2 --> GPU2
        GY1 -- 输入 --> GPU1
        GY2 -- 输入 --> GPU2
        GPU1 -- 计算 GY1 @ B1 --> Z1
        GPU2 -- 计算 GY2 @ B2 --> Z2
        Z1 & Z2 -->|All-Reduce (求和)| Z[最终输出 Z]
        Z -->|分发到所有GPU| GPU1_Z & GPU2_Z
    end

优点：

• 解决了单层过大的问题： 能够将巨大的权重矩阵分散到多个 GPU，有效降低了单个 GPU 的显存峰值。
• 计算与通信可以重叠： 在某些情况下，计算和通信（如 All-Reduce）可以并行进行，隐藏部分通信延迟。
• 相对较低的通信频率： 通信只在被切分的层之间发生，而不是每个 mini-batch 之后都进行全局梯度同步。

缺点：

• 实现复杂： 需要对模型的前向和反向传播逻辑进行深入修改，手动处理张量的切分、计算和聚合。
• 通信开销较高：All-Reduce 的通信量与模型的隐藏层维度和序列长度有关，当并行度增加时，通信开销会变得显著。
• 不适用于所有层： 主要适用于具有大权重矩阵的层，如线性层和注意力机制。对于像 LayerNorm 这样难以切分的操作，效果不佳。

简单的 PyTorch TP 代码示例 (概念性)：

实现完整的 TP 非常复杂，这里提供一个概念性的伪代码，帮助理解其核心逻辑。

import torch
import torch.nn as nn
import torch.distributed as dist

# 假设已在2个GPU上初始化了分布式环境
# rank = dist.get_rank()
# world_size = dist.get_world_size()

class ColumnParallelLinear(nn.Module):
    def __init__(self, input_size, output_size, world_size):
        super().__init__()
        self.input_size = input_size
        self.output_size_per_partition = output_size // world_size
        self.world_size = world_size

        # 每个GPU只创建部分权重的切片
        self.weight = nn.Parameter(torch.randn(self.input_size, self.output_size_per_partition))

    def forward(self, x):
        # 输入x在所有GPU上是完整的
        # 计算部分结果
        partial_output = torch.matmul(x, self.weight)
        return partial_output

class RowParallelLinear(nn.Module):
    def __init__(self, input_size, output_size, world_size):
        super().__init__()
        self.input_size_per_partition = input_size // world_size
        self.output_size = output_size
        self.world_size = world_size

        # 每个GPU只创建部分权重的切片
        self.weight = nn.Parameter(torch.randn(self.input_size_per_partition, self.output_size))

    def forward(self, x):
        # 输入x是按列切分的，每个GPU只有一部分
        # 计算部分结果
        partial_output = torch.matmul(x, self.weight)
        # 关键步骤：使用All-Reduce聚合结果
        dist.all_reduce(partial_output)
        return partial_output

# 概念性模型
class SimpleTPModel(nn.Module):
    def __init__(self, world_size):
        super().__init__()
        # 第一个线性层按列并行
        self.linear1 = ColumnParallelLinear(10, 20, world_size)
        # 第二个线性层按行并行
        self.linear2 = RowParallelLinear(20, 5, world_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        # x -> [batch_size, 10]
        output1 = self.linear1(x) # output1 -> [batch_size, 20 / world_size]
        output1 = self.relu(output1)
        output2 = self.linear2(output1) # output2 -> [batch_size, 5] (经过All-Reduce后)
        return output2

# 实际使用中，需要处理输入的广播、梯度的同步等更多细节
# 像Megatron-LM, DeepSpeed, PyTorch 2.x的DTensor等库封装了这些复杂性
print("这是一个TP的概念性代码示例，展示了列并行和行并行的核心思想。")

3. 流水线并行 (Pipeline Parallelism, PP)：像工厂流水线一样处理模型

核心思想： 流水线并行将模型的不同层（或模块）分配到不同的 GPU 上，形成一个“流水线”。数据像在工厂流水线上一样，依次流过每个 GPU，完成一部分计算。它的核心思想是“层间并行”。

工作流程：

1. 模型切分： 将一个大的模型（如一个包含 48 层的 Transformer）按层切分成多个阶段（stage）。例如，可以切分成 4 个 stage，每个 stage 包含 12 层，分别放置在 4 个 GPU 上。
2. 微批次 (Micro-batch)： 为了提高 GPU 的利用率，将一个大的 batch 进一步切分成多个更小的“微批次”（micro-batch）。
3. 流水线执行：

• GPU 1 (Stage 1) 处理第一个微批次（MB1），计算完成后将其激活值（activations）发送给 GPU 2。
• 在 GPU 2 (Stage 2) 开始处理 MB1 的同时，GPU 1 可以立即开始处理第二个微批次（MB2）。
• 以此类推，数据在流水线中流动，理想情况下，所有 GPU 都在同时处理不同的微批次。

流水线气泡 (Pipeline Bubble)：

流水线并行的一个主要问题是“气泡”——即 GPU 的空闲时间。在流水线的启动（warm-up）和排空（cool-down）阶段，部分 GPU 会处于等待状态，无法满负荷工作。微批次切分的越细，气泡带来的额外开销就越小，但同时也会增加通信次数。

图解PP：

gantt
    title Pipeline Parallelism (4 Stages, 4 Micro-batches)
    dateFormat  X
    axisFormat  %L
    section GPU 1 (Stage 1)
    Fwd MB1: 0, 1
    Fwd MB2: 1, 1
    Fwd MB3: 2, 1
    Fwd MB4: 3, 1
    Bwd MB1: 6, 1
    Bwd MB2: 7, 1
    Bwd MB3: 8, 1
    Bwd MB4: 9, 1
    section GPU 2 (Stage 2)
    Fwd MB1: 1, 1
    Fwd MB2: 2, 1
    Fwd MB3: 3, 1
    Fwd MB4: 4, 1
    Bwd MB1: 5, 1
    Bwd MB2: 6, 1
    Bwd MB3: 7, 1
    Bwd MB4: 8, 1
    section GPU 3 (Stage 3)
    Fwd MB1: 2, 1
    Fwd MB2: 3, 1
    Fwd MB3: 4, 1
    Fwd MB4: 5, 1
    Bwd MB1: 4, 1
    Bwd MB2: 5, 1
    Bwd MB3: 6, 1
    Bwd MB4: 7, 1
    section GPU 4 (Stage 4)
    Fwd MB1: 3, 1
    Fwd MB2: 4, 1
    Fwd MB3: 5, 1
    Fwd MB4: 6, 1
    Bwd MB1: 3, 1
    Bwd MB2: 4, 1
    Bwd MB3: 5, 1
    Bwd MB4: 6, 1

（上图展示了前向传播（Fwd）和反向传播（Bwd）在4个GPU上的执行情况，空白部分即为“流水线气泡”）

优点：

• 显著降低单卡显存： 每个 GPU 只需存储模型的一部分层，因此可以训练远超单卡显存容量的巨大模型。
• 通信开销相对较低： 通信只发生在相邻的 stage 之间，通信量是每个微批次的激活值大小，相比于 TP 的 All-Reduce，通常通信量更小。

缺点：

• 流水线气泡导致效率下降： GPU 的空闲时间会降低整体的计算效率。
• 负载均衡困难： 将模型切分成计算量和参数量都相等的 stage 是一个难题。如果切分不均，会导致某些 stage 成为瓶颈。
• 实现复杂： 需要精细地管理微批次的调度、数据传输和梯度同步。

简单的 PyTorch PP 代码示例 (使用 torch.distributed.pipeline)

PyTorch 提供了 torch.distributed.pipeline 模块来简化流水线并行的实现。

import torch
import torch.nn as nn
from torch.distributed.pipeline.sync import Pipe
from torch.distributed import rpc

# 假设已在2个GPU上初始化了RPC环境
# rpc.init_rpc(...)

# 1. 定义模型的各个部分
model_part1 = nn.Sequential(
    nn.Linear(100, 200),
    nn.ReLU()
).cuda(0)

model_part2 = nn.Sequential(
    nn.Linear(200, 50),
    nn.ReLU()
).cuda(1)

# 2. 将模型部分组合成一个序列
model = nn.Sequential(model_part1, model_part2)

# 3. 关键步骤：使用Pipe包装模型
# chunks参数指定了微批次数量
# Pipe会自动处理模型的切分和跨设备通信
pipe_model = Pipe(model, chunks=8)

# 4. 模拟训练
# 在 rank 0 上执行
# if rank == 0:
#     inputs = torch.randn(16, 100).cuda(0)
#     outputs = pipe_model(inputs) # 前向传播
#     outputs.sum().backward()     # 反向传播

# pipe_model.step() # 参数更新

print("这是一个使用PyTorch Pipe的PP简化示例。")
# 实际运行需要完整的RPC设置和启动脚本。

4. 专家并行 (Expert Parallelism, EP)：稀疏激活的MoE模型专属策略

核心思想： 专家并行是专门为混合专家模型（Mixture of Experts, MoE）设计的并行策略。MoE 模型并非在每次计算时都激活所有参数，而是通过一个“路由器”（gating network）为每个输入 token 选择性地激活一小部分“专家”（expert，通常是前馈网络 FFN）。专家并行的核心思想是“专家分散，数据路由”。

工作流程：

1. 专家分布： 将模型中的大量专家（例如 64 个）分布到多个 GPU 上。例如，如果有 8 个 GPU，每个 GPU 可以承载 8 个专家。模型的其他部分（如自注意力层）通常在所有 GPU 上复制。
2. 路由器计算： 每个 GPU 独立计算其上的输入 token 应该由哪些专家来处理。路由器会为每个 token 生成一个权重分布，选择 top-k（通常 k=1 或 2）个专家。
3. 数据路由 (All-to-All)： 这是 EP 的核心和瓶颈。每个 GPU 需要将 token 发送到拥有其被分配专家的 GPU 上。这个过程通过一次 All-to-All 通信操作完成。例如，GPU 0 上的某些 token 可能需要由 GPU 3 上的专家处理，GPU 0 就需要将这些 token 的数据发送给 GPU 3。
4. 专家计算： 每个 GPU 在接收到从其他 GPU 路由过来的 token 后，用自己本地的专家对这些 token 进行计算。
5. 结果返回 (All-to-All)： 计算完成后，再通过一次 All-to-All 操作，将处理后的 token 结果发送回其原始的 GPU。

图解EP：

graph TD
    subgraph "输入 & 路由"
        T1[GPU1: Tokens] --> R1{Router 1}
        T2[GPU2: Tokens] --> R2{Router 2}
        T3[GPU...: Tokens] --> R3{Router ...}
        R1 -->|决定Token去向| D1[Dispatch Info 1]
        R2 -->|决定Token去向| D2[Dispatch Info 2]
        R3 -->|决定Token去向| D3[Dispatch Info ...]
    end
    subgraph "数据交换 (All-to-All)"
        D1 & D2 & D3 -->|All-to-All| Comm[Token Shuffle]
        Comm --> GPU1_Tokens & GPU2_Tokens & GPU3_Tokens
    end
    subgraph "专家计算 & 结果返回"
        GPU1_Tokens --> E1[GPU1: Experts 1-8]
        GPU2_Tokens --> E2[GPU2: Experts 9-16]
        GPU3_Tokens --> E3[GPU...: Experts ...]
        E1 --> Res1[Results 1]
        E2 --> Res2[Results 2]
        E3 --> Res3[Results ...]
        Res1 & Res2 & Res3 -->|All-to-All (Combine)| Final[Final Outputs on original GPUs]
    end

优点：

• 极大扩展模型规模： 可以在计算成本仅少量增加的情况下，将模型参数量扩展数倍甚至数十倍，因为每次前向传播只激活了一小部分参数。
• 计算效率高： 在专家计算阶段，是纯粹的本地计算，没有通信。

缺点：

• All-to-All 通信瓶颈：All-to-All 操作的通信量非常大，随着 GPU 数量的增加，很容易成为性能瓶颈。
• 负载不均： 路由器的分配策略可能导致某些专家接收到的 token 远多于其他专家，造成严重的负载不均衡问题。
• 仅适用于 MoE 架构： 这是一种特定于模型架构的并行策略。

简单的 PyTorch MoE + EP 代码示例 (使用 DeepSpeed)：

DeepSpeed 库对 MoE 和 EP 提供了良好的支持。

import torch
import torch.nn as nn
# import deepspeed

# 这是一个概念性的示例，展示了DeepSpeed中MoE层的定义
# 实际运行需要完整的DeepSpeed环境和配置

class MyExpert(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.layer = nn.Linear(in_features, out_features)
    def forward(self, x):
        return self.layer(x)

# 假设我们有8个专家，并且希望将它们分布在2个GPU上（ep_size=2）
# 这意味着每个GPU将持有4个专家
# num_experts = 8
# ep_size = 2

# class MyMoEModel(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.embedding = nn.Embedding(1000, 512)
#         # 关键步骤：定义MoE层
#         # DeepSpeed会处理专家的分布和All-to-All通信
#         self.moe_layer = deepspeed.moe.layer.MoE(
#             hidden_size=512,
#             expert=MyExpert(512, 512),
#             num_experts=num_experts,
#             ep_size=ep_size,
#             k=1 # 每个token选择1个专家
#         )
#         self.output_layer = nn.Linear(512, 10)

#     def forward(self, x):
#         x = self.embedding(x)
#         # MoE层会返回输出和辅助损失（用于负载均衡）
#         output, _, _ = self.moe_layer(x)
#         output = self.output_layer(output)
#         return output

print("这是一个使用DeepSpeed实现MoE和EP的概念性代码示例。")

5. 序列并行 (Sequence Parallelism, SP)：解决长序列的显存难题

核心思想： 在处理超长序列（如长文档、高分辨率图像）时，注意力机制中的 (Query @ Key.T) 矩阵会变得异常巨大，其大小与序列长度的平方成正比，极易导致显存溢出。序列并行旨在解决这个问题，其核心思想是“序列切分，逐块计算”。

工作流程（以自注意力为例）：

1. 序列切分： 将输入的长序列（例如长度为 8192）在序列维度上进行切分，分发到不同的 GPU。例如，使用 4 个 GPU，每个 GPU 获得一个长度为 2048 的子序列。
2. 局部计算： 在不需要跨 GPU 通信的层（如 LayerNorm、MLP 层）上，每个 GPU 可以独立地对自己的子序列进行计算。这一步是和张量并行（TP）结合的。
3. 分布式注意力计算： 这是 SP 的核心。计算注意力时，需要完整的 Query, Key, Value 矩阵。

• 首先，每个 GPU 根据自己的子序列计算出局部的 Q_i, K_i, V_i。
• 然后，通过 All-Gather 操作，每个 GPU 都收集到所有其他 GPU 的 K 和 V，从而在本地拼凑出完整的 K 和 V 矩阵。
• 现在，每个 GPU 拥有局部的 Q_i 和完整的 K, V。它可以计算出完整的注意力分数 S_i = Q_i @ K.T 和部分的输出 O_i = S_i @ V。
• 在反向传播时，也需要类似的操作来计算梯度。

与TP的区别：

• TP切分隐藏层维度，SP切分序列维度。
• 在 Transformer 中，TP 主要用于切分 MLP 层的权重，而 SP 主要用于处理自注意力层中的序列维度。
• SP 通常与 TP 结合使用，以在不同层上实现最优的并行策略。

优点：

• 有效处理超长序列： 直接解决了由序列长度过长导致的显存瓶颈问题。
• 与 TP 互补： 可以与 TP 结合，对 Transformer 的不同部分采用最合适的并行策略。

缺点：

• 通信开销大：All-Gather 操作的通信量与序列长度、隐藏层维度和 GPU 数量都相关，开销不菲。
• 实现非常复杂： 需要对注意力机制的计算流程进行深度重构。

简单的 PyTorch SP 代码示例 (概念性)：

完整的 SP 实现非常复杂，通常集成在像 Megatron-LM 或 Colossal-AI 这样的框架中。以下是一个高度简化的概念伪代码。

import torch
import torch.nn as nn
import torch.distributed as dist

# 假设已在2个GPU上初始化了分布式环境
# rank = dist.get_rank()
# world_size = dist.get_world_size()
# seq_len = 8192
# hidden_size = 512
# local_seq_len = seq_len // world_size

# 1. 输入数据按序列维度切分
# input_x: [batch, seq_len, hidden_size]
# local_input_x = input_x.split(local_seq_len, dim=1)[rank] # [batch, local_seq_len, hidden_size]

class SequenceParallelAttention(nn.Module):
    def __init__(self):
        super().__init__()
        # Q, K, V的投影层权重在所有GPU上是完整的（或经过TP切分）
        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, hidden_size)
        self.v_proj = nn.Linear(hidden_size, hidden_size)

    def forward(self, local_x):
        # local_x: [batch, local_seq_len, hidden_size]
        
        # 计算局部的Q, K, V
        local_q = self.q_proj(local_x) # [batch, local_seq_len, hidden_size]
        local_k = self.k_proj(local_x) # [batch, local_seq_len, hidden_size]
        local_v = self.v_proj(local_x) # [batch, local_seq_len, hidden_size]

        # 关键步骤：All-Gather K和V
        # 创建一个列表来存放来自所有GPU的K
        k_list = [torch.empty_like(local_k) for _ in range(world_size)]
        dist.all_gather(k_list, local_k)
        full_k = torch.cat(k_list, dim=1) # [batch, seq_len, hidden_size]

        v_list = [torch.empty_like(local_v) for _ in range(world_size)]
        dist.all_gather(v_list, local_v)
        full_v = torch.cat(v_list, dim=1) # [batch, seq_len, hidden_size]

        # 计算注意力
        # local_q @ full_k.T
        attn_scores = torch.matmul(local_q, full_k.transpose(-1, -2))
        attn_probs = torch.softmax(attn_scores, dim=-1)

        # local_attn_output
        local_output = torch.matmul(attn_probs, full_v) # [batch, local_seq_len, hidden_size]
        
        return local_output

print("这是一个SP的概念性代码示例，展示了如何通过All-Gather实现分布式注意力计算。")

6. 五种并行策略的对比与总结

为了更清晰地理解这五种策略，我们从多个维度进行对比：

如何选择与组合？

在实践中，单一的并行策略往往无法满足需求，通常需要将它们组合使用，形成混合并行策略。

• TP + PP： 这是最常见的组合。当模型既深（层数多）又宽（隐藏层维度大）时，可以先用 PP 将模型按层切分到不同节点（服务器），然后在每个节点内部，再用 TP 将单个大层切分到该节点的所有 GPU 上。
• DP + TP + PP： 在 TP+PP 的基础上，还可以引入 DP。此时，整个 TP+PP 构成一个“虚拟的”单一设备，然后对这个“虚拟设备”进行数据并行。也就是说，我们可以复制多份 TP+PP 的模型副本，每份副本处理不同的数据。这被称为 3D 并行。
• SP 的角色： 当以上组合仍然因为序列长度问题而遭遇瓶颈时，就需要引入 SP。SP 通常与 TP 结合，在 Transformer 层级进行优化，用 SP 优化自注意力层，用 TP 优化 MLP 层。
• EP 的角色： EP 是 MoE 模型的标配。当使用 MoE 架构时，EP 是必须考虑的并行维度，它通常和其他并行策略（如 TP、DP）结合，以实现对巨大稀疏模型的有效训练。

选择的通用思路：

1. 首先考虑数据并行 (DP)： 如果模型可以完整地放入单张 GPU，但希望通过增大 batch size 来加速训练，DP 是最简单有效的选择。
2. 模型太大，单卡放不下？

• 如果是个别层（如 MLP）特别大，优先考虑 **张量并行 (TP)**。
• 如果是整个模型层数非常多，导致整体放不下，优先考虑 **流水线并行 (PP)**。
• 在大多数情况下，TP 和 PP 会组合使用。

结语

DP、EP、TP、PP、SP 这五种并行策略，是从不同维度对大规模模型训练这一复杂问题给出的解决方案。它们各有侧重，也各有局限。数据并行胜在简单普适，是提升吞吐量的基础；张量并行和流水线并行是解决模型“过大”而无法装入单卡的核心武器，前者治“胖”，后者治“高”；序列并行则专注于攻克“过长”序列的难题；而专家并行则为稀疏化模型打开了通向更大参数规模的大门。

在通往通用人工智能的道路上，模型规模的增长仍未停歇。深刻理解并熟练运用这些并行技术，将这些“积木”灵活地组合与创新，是每一位大模型算法工程师和系统工程师的必备技能。希望这篇详尽的解析，能为您在探索大模型的征途上，提供一份清晰的地图和有力的工具。