别再单卡硬扛了：一文讲透 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 原理
- 数据切分
- 多进程模型
其实就没那么难了。
