PyTorch分布式训练怎么调优 - GPU

PyTorch 分布式训练（DDP / FSDP / 单机多卡 / 多机多卡）的“调优”通常不是单点优化，而是系统性工程。下面按「优先级 + 可落地做法」给你一份实用调优清单。

一、先确认瓶颈在哪（最重要）

不要盲目调参，先定位：

GPU 利用率低（nvidia-smi 波动大）→ 数据 / 通信瓶颈
GPU 满但训练慢 → 计算 / 通信效率低
OOM → 显存优化
多机明显变慢 → 网络 / 通信后端问题

常用工具：

nvidia-smi dmon
torch.profiler
py-spy / nsight

二、分布式模式选择（基础）

场景	推荐
单机多卡	DDP（默认最优）
大模型 / 显存不够	FSDP
超大模型	FSDP + CPU offload
推理	DDP / Tensor Parallel

优先用 DDP，不是 DataParallel

三、DDP 关键调优项（最常踩坑）

1. 正确初始化

dist.init_process_group(
backend="nccl",
init_method="env://"
)

GPU → 必须用 nccl
CPU → gloo

2. 数据加载是头号瓶颈

✅ 正确做法：

train_sampler = DistributedSampler(dataset)
loader = DataLoader(
dataset,
batch_size=...,
num_workers=4~8,
pin_memory=True,
prefetch_factor=2
)

❌ 常见错误：

num_workers=0
每卡 batch 太小（<16）
数据解码在主进程

✅ 提升效果明显：

用 WebDataset / TFRecord
图像预处理放 GPU（可选）

3. 通信优化（NCCL）

环境变量（非常关键）：

export NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_P2P_LEVEL=NVL
export NCCL_IB_DISABLE=0   # 多机用 IB

多机：

用 InfiniBand
避免 tcp

4. 梯度通信与同步

DDP 自动 AllReduce，但可以优化：
✅ 使用：

torch.cuda.amp.autocast()
torch.cuda.amp.GradScaler()

✅ 梯度压缩（极端场景）：

torch.distributed.algorithms.ddp_comm_hooks

四、Batch Size 与 LR 调优

线性缩放规则

new_lr = base_lr * (world_size * batch_per_gpu / base_batch)

⚠ 注意：

1024 总 batch 通常需要 LR warmup
warmup 步数 ∝ batch size

五、显存优化（OOM 必看）

1. 混合精度（必开）

with torch.cuda.amp.autocast():
output = model(input)

2. Gradient Checkpointing

model.gradient_checkpointing_enable()

3. FSDP（大模型）

from torch.distributed.fsdp import FSDP

选项：

sharding_strategy=FULL_SHARD
cpu_offload

六、多机训练调优

网络

28G+ 带宽
NCCL socket 正确绑定：

export NCCL_SOCKET_IFNAME=eth0

启动方式

✅ 推荐：

torchrun --nproc_per_node=8 --nnodes=2 train.py

七、常见性能损失点（ checklist ）

[ ] DDP 没用 DistributedSampler
[ ] num_workers 太少
[ ] 每卡 batch 太小
[ ] 没开 AMP
[ ] print / logging 在主进程以外
[ ] 多机走 TCP 而不是 IB
[ ] optimizer.step 在 DDP 外

八、一个“标准高性能模板”

model = model.to(device)
model = DDP(model, device_ids=[local_rank])
optimizer = ...
for x, y in loader:
x, y = x.cuda(), y.cuda()
with torch.cuda.amp.autocast():
loss = model(x, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

九、如果你愿意，我可以进一步帮你

你可以告诉我：

单机还是多机？
GPU 数量和型号？
模型类型（CNN / Transformer / LLM）
当前 GPU 利用率大概多少？

我可以直接给你针对性配置和代码片段。