别再单卡硬扛了：一文讲透 Python 多 GPU / 分布式训练怎么写（附完整实战代码）

大家好，我是 Echo_Wish。

说句实话，我第一次接触多 GPU 训练的时候，内心是崩溃的。

当时的我还停留在：

👉 model.cuda() 就完事了

结果一上服务器，看到 8 张 GPU 闪闪发光，我却只用了一张——
那种感觉就像你租了 8 栋别墅，结果只睡厕所。

所以今天这篇文章，我不讲虚的，就带你从单卡 → 多卡 → 分布式，一步一步把这事讲明白，而且保证你能跑起来。

一、为什么你必须学多 GPU？

先别急着写代码，先搞清楚一件事：

👉 多 GPU 不只是“快”，而是“能不能跑”的问题

比如：

• 大模型（参数爆炸）
• 大 batch（稳定训练）
• 大数据（吞吐压力）

如果你只用单卡：

👉 要么 OOM
👉 要么训练 3 天

二、最简单的多 GPU：DataParallel（不推荐但好理解）

先从最容易上手的开始。

import torch
import torch.nn as nn

model = MyModel()
model = nn.DataParallel(model)
model = model.cuda()

训练代码不用改。

它是怎么工作的？

👉 一句话：

把 batch 切开 → 分发到多个 GPU → 汇总梯度

但问题也很明显：

• 主卡（GPU0）压力巨大
• 通信效率低
• 性能一般

👉 所以：

DataParallel 只适合入门，不适合生产

三、主流方案：DistributedDataParallel（DDP）

真正该用的是这个：

👉 DistributedDataParallel

核心思想（你一定要理解）

👉 每个 GPU 一个进程（而不是一个线程）

这点非常关键。

训练结构图（帮助你理解）

你可以这样理解：

• 每个 GPU：

  • 有自己模型副本
  • 处理自己数据

• 每一步：

  • 梯度同步（AllReduce）

四、DDP 最小可运行代码（强烈建议收藏）

1️⃣ 初始化环境

import torch.distributed as dist

def setup(rank, world_size):
    dist.init_process_group(
        backend="nccl",
        rank=rank,
        world_size=world_size
    )

2️⃣ 包装模型

from torch.nn.parallel import DistributedDataParallel as DDP

model = MyModel().to(rank)
model = DDP(model, device_ids=[rank])

3️⃣ 使用 DistributedSampler（重点！）

from torch.utils.data.distributed import DistributedSampler

train_sampler = DistributedSampler(dataset)

train_loader = DataLoader(
    dataset,
    batch_size=32,
    sampler=train_sampler
)

👉 为什么要这个？

👉 避免不同 GPU 读到同样数据

4️⃣ 训练循环

for epoch in range(epochs):
    train_sampler.set_epoch(epoch)
    
    for data, label in train_loader:
        data = data.to(rank)
        label = label.to(rank)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()

5️⃣ 启动方式（关键）

torchrun --nproc_per_node=4 train.py

👉 这句话的意思：

启动 4 个进程 = 4 张 GPU

五、很多人踩的坑（我帮你踩过了）

❌ 坑 1：忘了用 DistributedSampler

结果：

👉 每张卡都在训练同一批数据

= 白跑

❌ 坑 2：没有设置 device

torch.cuda.set_device(rank)

不然：

👉 GPU 会乱用

❌ 坑 3：打印日志混乱

解决：

if rank == 0:
    print("只让主进程输出")

❌ 坑 4：保存模型出错

if rank == 0:
    torch.save(model.state_dict(), "model.pth")

六、再进阶一点：多机分布式（跨服务器）

如果你有多台机器：

👉 本质没变，只是多了网络通信

关键参数

torchrun \
  --nnodes=2 \
  --nproc_per_node=4 \
  --node_rank=0 \
  --master_addr="192.168.1.1" \
  --master_port=29500 \
  train.py

理解一下：

• nnodes：机器数
• nproc_per_node：每台 GPU 数
• master_addr：主节点

七、性能优化（真正拉开差距的地方）

✅ 1. 混合精度训练（必开）

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    output = model(data)
    loss = criterion(output, label)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

👉 效果：

• 更快
• 更省显存

✅ 2. 梯度累积（显存不够用）

loss = loss / accumulation_steps
loss.backward()

✅ 3. 合理 batch size

经验：

👉 GPU 越多，batch 要跟着放大

八、我自己的一点真实感受

说点实话，多 GPU / 分布式这块，很多人卡在两个点：

1️⃣ “看懂了，但跑不起来”

原因很简单：

👉 环境问题 + 启动方式

2️⃣ “跑起来了，但不快”

原因：

👉 通信瓶颈

所以你要记住一句话：

👉 分布式训练，本质不是“算力问题”，而是“通信问题”

九、什么时候该用？什么时候别用？

我给你一个非常实用的判断标准：

✅ 用多 GPU：

• 模型大（比如 Transformer）
• 数据多
• 单卡训练慢

❌ 别用：

• 小模型（反而更慢）
• 调试阶段（会崩溃你心态）

十、最后总结一句话（重点）

如果你今天只记住一句话，那就是：

👉 DataParallel 是玩具，DDP 才是生产力

写在最后

我一直觉得，多 GPU 训练这件事，本质上不是“技术门槛高”，而是：

👉 信息太碎 + 坑太多

你一旦把这几个关键点搞懂：

• DDP 原理
• 数据切分
• 多进程模型

其实就没那么难了。