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

模型并行是将单个模型拆分到多卡，各卡负责不同层；而DataParallel和DDP是数据并行，需模型完整副本在每卡。仅当单卡显存不足时才用模型并行，否则优先选DDP。

如何在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')
<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同理，它要求模型所有参数位于同一设备。

python-flask-vue-creator

创建 Flask + Vue.js 全栈项目，采用单端口架构、uv 依赖管理、结构化日志和生产级前端设计。当用户要求以下操作时触发：创建新的 Flask+Vue 项目、搭建 Web 应用、构建全栈 Python+Vue 应用、初始化 Flask Web 项目、设置单端口 Flask 前后端项目。触发词：新建项目、创建项目、搭建项目、Flask项目、Vue项目。

下载

常见误操作：

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