当前位置：首页 > 文章列表 > 文章 > python教程 > PyTorch模型并行训练详解

PyTorch模型并行训练详解

2026-05-07 23:35:03 0浏览收藏

本文深入解析了PyTorch中易被混淆的“模型并行”本质——它并非DataParallel或DDP的数据并行变体，而是将单一大模型按层或模块显式拆分到多张GPU上以突破单卡显存限制的精细策略；文章直击核心：仅当模型大到无法装入单卡（如超大语言模型）时才应选用模型并行，否则DDP更高效可靠，并手把手揭示了手动设备分配的关键细节——从模块级.to(device)初始化、forward中严谨的跨卡张量搬运，到避免常见设备错位错误和PCIe带宽陷阱，还点明了复杂场景下借助DeepSpeed等工具的必要性，为面临显存瓶颈的开发者提供了清晰、务实且避坑指南。

如何在Python中实现PyTorch的模型并行训练_将不同层分配至不同GPU

PyTorch模型并行（Model Parallel）不是DataParallel或DDP

模型并行和数据并行是两回事：torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel 都是把**同个模型副本**复制到多个GPU上，靠分batch做数据并行；而模型并行是把**单个模型的不同子模块（比如前几层放GPU0，后几层放GPU1）** 拆开部署，显存压力分散，但引入跨GPU张量传输开销。

关键判断：只有当单卡显存塞不下整个模型（例如大语言模型、超深CNN）时才值得用模型并行；否则优先用DistributedDataParallel——它更简单、通信更高效、框架支持更完善。

手动将不同`nn.Module`分配到不同GPU的实操步骤

核心就是显式调用.to(device)，并在forward中处理设备间数据搬运。PyTorch不自动帮你搬，漏了就会报Expected all tensors to be on the same device错误。

定义模型时，按逻辑切分成多个子模块（如self.encoder、self.decoder），每个子模块单独.to('cuda:0')或.to('cuda:1')
forward里，确保输入先送到第一块GPU，中间结果在传给下一块GPU前显式调用.to('cuda:1')
损失计算和backward()前，所有参与运算的张量必须在同一设备上（通常统一挪回cuda:0或就地在最后一块GPU上算）
优化器只对模型参数起作用，不管参数在哪个GPU——只要参数本身绑定了设备，optimizer.step()就能正常更新

示例片段：

class ModelParallelMLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.layer1 = nn.Linear(input_size, hidden_size).to('cuda:0')
        self.layer2 = nn.Linear(hidden_size, output_size).to('cuda:1')
def forward(self, x):
    x = x.to('cuda:0')           # 输入先上cuda:0
    x = self.layer1(x)
    x = x.to('cuda:1')          # 中间结果搬去cuda:1
    x = self.layer2(x)
    return x

为什么不能直接用`nn.DataParallel`或`DDP`做模型并行

DataParallel会尝试把整个模型复制到每张卡，遇到部分参数在cuda:0、部分在cuda:1时直接报错：Module 'XXX' has parameters on different devices；DistributedDataParallel同理，它要求模型所有参数位于同一设备。

常见误操作：

在__init__里混合调用.to('cuda:0')和.to('cuda:1')，但没在forward里做设备同步 → 运行时报device mismatch
用了torch.cuda.set_device(1)全局设卡，结果layer1被意外挪到cuda:1，和预设冲突
想用torch.nn.parallel.scatter或gather来“辅助”模型并行——它们是为DataParallel内部服务的，不适用于手动模型拆分