用PyTorch在超大规模下训练深度学习模型：并行策略全解析

我猜咱们每个人肯定都累坏了，天天追着 LLM 研究社区跑，感觉每天都冒出个新的最牛模型，把之前的基准都给打破了呢。要是你好奇为啥创新速度能这么快，那主要就是研究人员能够在超大规模下训练和验证模型啦，这全靠并行计算的功劳呀。

要是你还没听说过呢，5D 并行这个术语最早是 Meta AI 的论文 — The Llama 3 Herd of Models 里火起来的哦。传统上，它指的是结合数据、张量、上下文、流水线和专家并行这些技术呢。不过最近呀，又冒出了个新的玩意儿 — ZeRO (Zero Redundancy Optimizer)，这可是个大杀器，通过减少分布式计算中的冗余来优化内存呢。每种技术都针对训练挑战的不同方面，它们组合起来就能搞定那些有几十亿甚至几万亿参数的模型啦。

我之前已经讲过一些比较底层的技巧啦，这些技巧能让你用 PyTorch 训练和部署模型的速度更快哦。虽然这些小贴士和小技巧能在训练时给你加分，但它们也就是加分项而已呀。要是你没有从根本上搞懂它们该在啥时候、啥地方用，那很可能就会用错地方啦。

这篇文章的重点呢，就是要给大家讲清楚模型操作的高层组织结构，还会用 PyTorch 来举些例子哦。这些并行计算的基本原则，就是现在能让模型在超大规模下（想想看，日活跃用户有几千万呢）更快迭代和部署的关键因素哦。

1. 数据并行

数据并行 [图片来源]

数据并行是最简单也是最常用的并行技术啦。它就是创建多个相同的模型副本，然后让每个副本在数据的不同子集上进行训练呢。在本地计算完梯度后，就会通过全归约操作（all-reduce operation）把梯度聚合起来，用来更新所有模型副本的参数呢。

当模型本身能装进单个 GPU 的内存里，但数据集太大没法按顺序处理的时候，这种方法特别管用哦。

PyTorch 通过 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel(DDP) 模块，为数据并行提供了现成的支持呢。这其中呀，DDP 更受大家青睐，因为它在多节点设置下有更好的可扩展性和效率呢。NVIDIA 的 NeMo 框架就很好地展示了它是怎么工作的哦 —

数据并行示意图 [图片来源]

一个实现示例可能长这样：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义你的模型
model = nn.Linear(10, 1)

# 用 DataParallel 包裹模型
model = nn.DataParallel(model)

# 把模型移到 GPU
model = model.cuda()

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 假数据
inputs = torch.randn(64, 10).cuda()
targets = torch.randn(64, 1).cuda()

# 前向传播
outputs = model(inputs)
loss = criterion(outputs, targets)

# 反向传播和优化
loss.backward()
optimizer.step()

重点收获

• 小模型 / 大数据集 — 只有当模型能装进单个 GPU 的内存，但数据集太大时，这种方法才有效哦。
• 模型复制 — 每个 GPU 都会保存一份相同的模型参数副本呢。
• 小批量分割 — 输入数据会在 GPU 之间分配，确保每个设备处理一个独立的小批量数据哦。
• 梯度同步 — 在前向和反向传播之后，梯度会在 GPU 之间同步，以保持一致性呢。

优点和注意事项

• 简单高效 — 实现起来很简单，能很轻松地和现有的代码库集成，而且在处理大数据集时扩展性特别好哦。
• 通信开销 — 在梯度同步时的通信开销可能会成为大规模系统的一个瓶颈呢。

2. 张量并行

张量并行 [图片来源]

数据并行关注的是分割数据，而 张量并行（或者叫 模型并行）则是把模型本身分散到多个设备上呢。这种方法会把大型权重矩阵和中间张量分割开来，让每个设备只负责一部分计算呢。和数据并行不同（数据并行会在每个 GPU 上复制整个模型），张量并行会把模型的层或者张量分散到不同的设备上呢。每个设备负责计算模型前向和反向传播的一部分哦。

当模型太大，没办法装进单个 GPU 的内存时，这种方法就特别有用啦，尤其是对于那些基于 Transformer 的超大模型呢。

虽然 PyTorch 没有直接提供现成的张量并行支持，但用 PyTorch 灵活的张量操作和分布式通信原语，很容易就能实现自定义的张量并行呢。不过呢，要是想要更强大的解决方案，像 DeepSpeed 和 Megatron-LM 这样的框架就能扩展 PyTorch 来实现这个功能哦。一个简单的张量并行实现示例如下：

import torch
import torch.distributed as dist

def tensor_parallel_matmul(a, b, devices):
    # a 按行分割，b 在设备间共享
    a_shard = a.chunk(len(devices), dim=0)
    results = []
    for i, dev in enumerate(devices):
        a_device = a_shard[i].to(dev)
        b_device = b.to(dev)
        results.append(torch.matmul(a_device, b_device))
    # 把各个设备的结果拼接起来
    return torch.cat(results, dim=0)

# 示例用法：
a = torch.randn(1000, 512)  # 假设这个张量太大，一个 GPU 装不下
b = torch.randn(512, 256)
devices = ['cuda:0', 'cuda:1']

result = tensor_parallel_matmul(a, b, devices)

重点收获

• 大模型 — 当模型太大，装不下单个 GPU 的内存时，这种方法特别有效哦。
• 分割权重 — 不是在每个设备上复制整个模型，张量并行会把模型的参数切片呢。
• 集体计算 — 前向和反向传播是在 GPU 之间集体完成的，需要精心协调，以确保张量的所有部分都被正确计算呢。
• 自定义操作 — 往往要用到专门的 CUDA 内核或者第三方库，才能高效地实现张量并行哦。

优点和注意事项

• 内存效率 — 通过分割大型张量，你就能训练那些超出单个设备内存的模型啦。它还能显著减少矩阵操作的延迟呢。
• 复杂性 — 设备之间的协调增加了额外的复杂性哦。当扩展到超过两个 GPU 时，开发者必须仔细管理同步呢。由于手动分割可能导致的负载不平衡，以及为了避免 GPU 空闲太久而需要进行的设备间通信，是这些实现中常见的问题呢。
• 框架增强 — 像 Megatron-LM 这样的工具已经为张量并行树立了标杆，而且很多这样的框架都能和 PyTorch 无缝集成呢。不过，集成并不总是那么顺利哦。

3. 上下文并行

上下文并行采用了一种不同的方法，它针对的是输入数据的上下文维度，尤其在基于序列的模型（比如 Transformer）中特别厉害呢。主要思想就是把长序列或者上下文信息分割开来，让不同的部分同时进行处理呢。这样能让模型在不超出内存或者计算能力的情况下，处理更长的上下文哦。当需要一起训练多个任务时，比如在多任务 NLP 模型中，这种方法就特别有用啦。

和张量并行类似，PyTorch 本身并没有原生支持上下文并行呢。不过，通过巧妙地重构数据，我们就能有效地管理长序列啦。想象一下，有一个 Transformer 模型需要处理长文本 —— 可以把序列分解成更小的片段，然后并行处理，最后再合并起来呢。

下面就是一个自定义 Transformer 块中上下文如何分割的示例哦。在这个示例中，这个块可能会并行处理长序列的不同片段，然后把输出合并起来进行最后的处理呢。

import torch
import torch.nn as nn

class ContextParallelTransformer(nn.Module):
    def __init__(self, d_model, nhead, context_size):
        super(ContextParallelTransformer, self).__init__()
        self.context_size = context_size
        self.transformer_layer = nn.TransformerEncoderLayer(
                                    d_model=d_model, nhead=nhead)

    def forward(self, x):
        # x 的形状：[batch, seq_len, d_model]
        batch, seq_len, d_model = x.size()
        assert seq_len % self.context_size == 0, \
                "序列长度必须能被 context_size 整除"
        # 把序列维度分割成片段
        segments = x.view(batch, seq_len // self.context_size,
                          self.context_size, d_model)
        # 使用循环或者并行映射并行处理每个片段
        processed_segments = []
        for i in range(segments.size(1)):
            segment = segments[:, i, :, :]
            processed_segment = self.transformer_layer(
                                segment.transpose(0, 1))
            processed_segments.append(processed_segment.transpose(0, 1))
        # 把处理过的片段拼接回完整的序列
        return torch.cat(processed_segments, dim=1)

# 示例用法：
model = ContextParallelTransformer(d_model=512, nhead=8, context_size=16)
# [batch, sequence_length, embedding_dim]
input_seq = torch.randn(32, 128, 512)
output = model(input_seq)

重点收获

• 序列分割 — 把序列或者上下文维度分割开来，就能在不同的数据片段上并行计算啦。
• 长序列的可扩展性 — 这对于处理特别长的序列的模型特别有用，要是把整个上下文一次性处理，那既不可能，也不高效哦。
• 注意力机制 — 在 Transformer 中，把注意力计算分割到不同的片段上，能让每个 GPU 处理序列的一部分以及它相关的自注意力计算呢。

优点和注意事项

• 高效的长序列处理 — 把长上下文分割成并行的片段，模型就能在不过度占用内存资源的情况下处理超长的序列啦。
• 序列依赖性 — 必须特别注意跨越上下文片段边界的依赖关系哦。可能需要采用重叠片段或者额外的聚合步骤等技术呢。
• 新兴领域 — 随着研究的不断深入，我们期待会有更多专门促进 PyTorch 中上下文并行的标准工具和库出现呢。

4. 流水线并行

流水线并行示意图 [图片来源]

流水线并行引入了把神经网络分割成一系列阶段的概念，每个阶段都在不同的 GPU 上进行处理呢。当数据流经网络时，中间结果会从一个阶段传递到下一个阶段，就像流水线一样呢。这种错开的执行方式能让计算和通信重叠起来，从而提高整体的吞吐量呢。

幸运的是，PyTorch 有一个现成的 API 支持这个功能，叫做 Pipe，用它就能非常轻松地创建分段的模型呢。这个 API 会自动把一个顺序模型分割成微批次，这些微批次会在指定的 GPU 上流动呢。

一个简单的使用示例如下：

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

# 定义模型的两个顺序片段
segment1 = nn.Sequential(
    nn.Linear(1024, 2048),
    nn.ReLU(),
    nn.Linear(2048, 2048)
)

segment2 = nn.Sequential(
    nn.Linear(2048, 2048),
    nn.ReLU(),
    nn.Linear(2048, 1024)
)

# 使用 Pipe 把片段组合起来
# 如果提供了设备分配，Pipe 会自动处理模块在设备上的放置
# 这里是自动分配的
model = nn.Sequential(segment1, segment2)
model = Pipe(model, chunks=4)

# 现在，当你把数据传递给模型时，微批次就会以流水线的方式进行处理啦。
inputs = torch.randn(16, 1024)
outputs = model(inputs)import torch
import torch.nn as nn
from torch.distributed.pipeline.sync import Pipe

# 定义模型片段
segment1 = nn.Sequential(
    nn.Linear(1024, 2048),
    nn.ReLU(),
    nn.Linear(2048, 2048)
)
segment2 = nn.Sequential(
    nn.Linear(2048, 2048),
    nn.ReLU(),
    nn.Linear(2048, 1024)
)

# 使用 Pipe 把片段组合成一个模型
model = nn.Sequential(segment1, segment2)
# 把模型分割成微批次，这些微批次会在 'cuda:0' 和 'cuda:1' 设备上流动
model = Pipe(model, devices=['cuda:0', 'cuda:1'], chunks=4)

# 模拟输入批次
inputs = torch.randn(16, 1024).to('cuda:0')
outputs = model(inputs)

重点收获

• 分阶段计算 — 把模型分割成一系列阶段（或者叫“流水线”）。每个阶段都分配给不同的 GPU 哦。
• 微批次 — 不是把一个大批次一次性传给一个阶段，而是把批次分割成微批次，这些微批次会持续不断地流经流水线呢。
• 提高吞吐量 — 通过确保所有设备都在同时工作（即使是在处理不同的微批次），流水线并行可以显著提高吞吐量哦。

优点和注意事项

• 资源利用 — 流水线并行可以通过重叠不同阶段的计算，提高 GPU 的利用率哦。
• 延迟与吞吐量的权衡 — 虽然吞吐量提高了，但可能会因为引入的流水线延迟，稍微影响一下延迟呢。
• 复杂的调度 — 有效的微批次调度和负载平衡对于在各个阶段实现最佳性能至关重要哦。

5. 专家并行

专家并行 [图片来源]

专家并行是一种受 混合专家模型（MoE） 启发的技术，旨在在保持计算成本可控的同时，扩展模型的容量呢。在这个范式中，模型由多个专门的“专家”组成 —— 这些子网络通过一个门控机制，为每个输入选择性地激活呢。对于每个样本，只有部分专家会参与处理，这样就能在不大幅增加计算开销的情况下，拥有巨大的模型容量啦。

混合专家模型使用的门控函数 [图片来源]

同样呢，PyTorch 并没有直接提供现成的专家并行解决方案哦，不过它模块化的设计使得创建自定义实现成为可能呢。这种策略通常涉及定义一组专家层，以及一个决定激活哪些专家的门控器呢。

在生产环境中，专家并行通常会和其他并行策略结合起来使用哦。比如，你可以同时使用数据并行和专家并行，既能处理大型数据集，又能处理大量的模型参数 —— 同时，还能通过门控机制，把计算有选择性地路由到合适的专家那里呢。下面就是一个简化版的实现示例：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Expert, self).__init__()
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return F.relu(self.fc(x))

class MoE(nn.Module):
    def __init__(self, input_dim, output_dim, num_experts, k=2):
        super(MoE, self).__init__()
        self.num_experts = num_experts
        self.k = k  # 每个样本使用的专家数量
        self.experts = nn.ModuleList([Expert(input_dim, output_dim)
                                      for _ in range(num_experts)])
        self.gate = nn.Linear(input_dim, num_experts)

    def forward(self, x):
        # x 的形状：[batch, input_dim]
        gate_scores = self.gate(x)  # [batch, num_experts]
        # 为每个输入选择 top-k 个专家
        topk = torch.topk(gate_scores, self.k, dim=1)[1]
        outputs = []
        for i in range(x.size(0)):
            expert_output = 0
            for idx in topk[i]:
                expert_output += self.experts[idx](x[i])
            outputs.append(expert_output / self.k)
        return torch.stack(outputs)

# 示例用法：
batch_size = 32
input_dim = 512
output_dim = 512
num_experts = 4
model = MoE(input_dim, output_dim, num_experts)
x = torch.randn(batch_size, input_dim)
output = model(x)

重点收获

• 混合专家 — 对于每个训练样本，只使用部分专家，这样就能在不大幅增加每个样本计算量的情况下，保持巨大的模型容量哦。
• 动态路由 — 门控函数会动态决定哪些专家应该处理每个输入标记或者数据片段呢。
• 专家级别的并行 — 专家可以在多个设备上分布开来，这样就能并行计算，进一步减少瓶颈啦。

优点和注意事项

• 可扩展的模型容量 — 专家并行让你能够构建出容量巨大的模型，而不会因为每个输入都增加计算量哦。
• 高效的计算 — 通过只为每个输入处理选定的专家子集，就能实现高效的计算啦。
• 路由复杂性 — 门控机制非常关键哦。要是设计得不好，可能会导致负载不平衡和训练不稳定呢。
• 研究前沿 — 专家并行仍然是一个活跃的研究领域，目前正在进行的研究旨在改进门控方法以及专家之间的同步呢。

6. ZeRO：零冗余优化器

分区策略和 GPU 性能 [图片来源]

ZeRO，也就是 零冗余优化器，在大规模训练的内存优化方面可是个突破性的成果哦。作为 DeepSpeed 库的一部分，ZeRO 通过分区优化器状态、梯度和模型参数，解决了分布式训练中的内存限制问题呢。说白了，ZeRO 就是消除了每个 GPU 都保存一份所有东西的冗余，从而节省了大量的内存呢。

它的运作方式是把优化器状态和梯度的存储分散到所有参与的设备上，而不是复制它们呢。这种策略不仅能减少内存使用量，还能让那些原本会超出单个 GPU 内存容量的模型也能进行训练呢。ZeRO 通常会分三个阶段来实现，每个阶段都针对不同的内存冗余问题：

ZeRO-1：优化器状态分区

• 把优化器状态（比如动量缓冲区）分区到各个 GPU 上呢
• 每个 GPU 只保存其参数部分的优化器状态哦
• 模型参数和梯度仍然在所有 GPU 上复制呢

ZeRO-2：梯度分区

• 包含了 ZeRO-1 的所有内容呢
• 另外还会把梯度分区到各个 GPU 上哦
• 每个 GPU 只计算并保存其参数部分的梯度呢
• 模型参数仍然在所有 GPU 上复制呢

ZeRO-3：参数分区

• 包含了 ZeRO-1 和 ZeRO-2 的所有内容呢
• 另外还会把模型参数分区到各个 GPU 上哦
• 每个 GPU 只保存模型参数的一部分呢
• 在前向和反向传播过程中需要收集参数呢
  

ZeRO Offload 的架构 [图片来源]

ZeRO 提供了最大的灵活性，因为它结合了数据和模型并行的好处，如上图所示呢。

虽然 ZeRO 是 DeepSpeed 的一个特性，但它和 PyTorch 的集成使得它成为了训练优化工具箱中的一个重要工具，有助于高效地管理内存，并让以前无法训练的模型大小在现代硬件上成为可能呢。下面是一个示例实现：

import torch
import torch.nn as nn
import deepspeed

class LargeModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LargeModel, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.relu(self.fc1(x))
        return self.fc2(x)

model = LargeModel(1024, 4096, 10)

# DeepSpeed 配置，带有 ZeRO 优化器设置
ds_config = {
    "train_batch_size": 32,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "zero_optimization": {
        "stage": 2,  # 第 2 阶段：梯度分区
        "allgather_partitions": True,
        "reduce_scatter": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True
    }
}

# 用 ZeRO 初始化 DeepSpeed 和模型
model_engine, optimizer, _, _ = deepspeed.initialize(model=model,
                                                     config=ds_config)
inputs = torch.randn(32, 1024).to(model_engine.local_rank)
outputs = model_engine(inputs)
loss = outputs.mean()  # 简化的损失计算
model_engine.backward(loss)
model_engine.step()

重点收获

• 阶段选择 — ZeRO 通常分多个阶段实现，每个阶段在内存节省和通信开销之间提供了不同的平衡呢。根据模型大小、网络能力以及可以接受的通信开销水平，选择合适的阶段至关重要哦。
• 与其他技术的集成 — 它可以无缝地融入一个可能还包括上述并行策略的生态系统中呢。

优点和注意事项

• 通信开销 — 这种策略的一个固有挑战是，减少内存冗余通常会增加 GPU 之间的数据交换量哦。因此，高效利用高速互连（比如 NVLink 或 InfiniBand）就变得更加关键啦，因为这个原因。
• 配置复杂性 — ZeRO 比传统优化器引入了更多的配置参数呢。这些设置需要仔细地进行实验和分析，以匹配硬件的优势，确保优化器能够高效运行呢。设置内容包括但不限于 — 适当的梯度聚合桶大小，以及各种状态（优化器状态、梯度、参数）的分区策略。
• 强大的监控 — 在启用 ZeRO 的训练中调试问题可能会非常困难哦。因此，提供对 GPU 内存使用情况、网络延迟以及整体吞吐量等信息的监控工具就变得至关重要啦，因为这个原因。

把它们全部结合起来

并行的各种混合方法 [图片来源]

在超大规模下训练深度学习模型，很多时候都需要采用混合方法 —— 通常会结合上述提到的这些技术呢。比如，一个最先进的 LLM 可能会使用数据并行来在节点之间分配批次，张量并行来分割巨大的权重矩阵，上下文并行来处理长序列，流水线并行来连接顺序模型阶段，专家并行来动态分配计算资源，最后再用 ZeRO 来优化内存使用呢。这种协同作用确保了即使是参数数量天文数字级别的模型，也仍然能够进行训练，并且保持高效的哦。

搞清楚在什么时候、在什么地方以及如何使用这些技术，对于突破可能的极限至关重要呢。再加上 PyTorch 的模块化和即插即用的库，构建能够突破传统硬件限制的健壮、可扩展的训练管道，已经变得越来越容易被更多人掌握了呢。