PyTorch模型并行训练详解

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

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')
<pre class="brush:python;toolbar:false;">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内部服务的，不适用于手动模型拆分

实际训练中容易忽略的性能陷阱

模型并行最大的代价不是计算，而是GPU间PCIe带宽瓶颈。比如在A100上，cuda:0→cuda:1传一个[32, 768]的float32张量，耗时可能高达0.1ms以上，远超单层计算时间。

• 尽量减少跨GPU张量传输次数：把能合并的层留在同一卡上，避免“一层一搬”
• 传输前用.contiguous() + .to()，避免因内存布局导致隐式拷贝放大延迟
• 不要在forward里写x = x.to('cuda:1'); x = x.to('cuda:0')来回搬——这等于白跑两次PCIe
• 验证时记得model.eval()后仍要保持设备一致，否则torch.no_grad()下也可能触发设备错位

真正复杂的模型并行（比如Transformer按注意力头/FFN切分）需要仔细设计通信粒度，这时候往往得结合torch.distributed.rpc或第三方库如DeepSpeed的pipeline_parallel，纯手工实现极易出错且难维护。

今天关于《PyTorch模型并行训练详解》的内容就介绍到这里了，是不是学起来一目了然！想要了解更多关于的内容请关注golang学习网公众号！