PyTorch分布式训练如何配置环境

下面系统介绍 PyTorch 分布式训练的环境配置，按「从零到能跑」的逻辑一步步来，适用于 单机多卡 / 多机多卡（DDP） 场景。

一、整体思路

PyTorch 分布式训练主流方案：

✅ 90% 场景用 DDP 即可

二、环境准备

1️⃣ 硬件 & 系统

• GPU ≥ 1 张（NVIDIA）
• Linux（Ubuntu 20.04/22.04 最稳）
• CUDA 驱动正常

nvidia-smi

2️⃣ 安装 CUDA & cuDNN

确认 CUDA 版本：

nvcc --version

建议：

• CUDA 11.8 / 12.1
• cuDNN ≥ 8.6

3️⃣ 安装 Python & PyTorch

推荐使用 conda

conda create -n torch_ddp python=3.10 -y
conda activate torch_ddp

安装 PyTorch（示例 CUDA 11.8）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证：

import torch
print(torch.cuda.device_count())
print(torch.cuda.is_available())

三、分布式核心组件

1️⃣ NCCL（必装）

✅ PyTorch 默认使用 NCCL 做 GPU 通信
一般 pip install torch 已包含，但服务器环境需确认：

apt install libnccl2 libnccl-dev

验证：

import torch
print(torch.distributed.is_nccl_available())

2️⃣ 环境变量（非常关键）

DDP 依赖以下环境变量：

四、单机多卡配置（最常见）

✅ 最简单方式（推荐）

使用 torchrun

torchrun \
--nproc_per_node=4 \
train.py

无需手动设置环境变量 ✅

✅ train.py 最小模板

import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
def setup():
dist.init_process_group(backend="nccl")
def cleanup():
dist.destroy_process_group()
def main():
setup()
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = nn.Linear(10, 1).cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
dataset = torch.randn(100, 10)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=8, sampler=sampler)
for x in loader:
x = x.cuda(local_rank)
loss = model(x).mean()
loss.backward()
cleanup()
if __name__ == "__main__":
main()

五、多机多卡配置

假设：

• 2 台机器
• 每台 4 张 GPU

节点 0（主节点）

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

节点 1

torchrun \
--nproc_per_node=4 \
--nnodes=2 \
--node_rank=1 \
--master_addr=192.168.1.10 \
--master_port=29500 \
train.py

⚠️ 多机必须：

• 网络互通
• 防火墙放行端口
• 所有机器代码 & 数据一致

六、Dataset 的正确写法（非常重要）

❌ 错误：

DataLoader(dataset, shuffle=True)

✅ 正确：

sampler = DistributedSampler(dataset)
DataLoader(dataset, batch_size=..., sampler=sampler)

每个 epoch 前：

sampler.set_epoch(epoch)

七、常见坑 & 排查

❌ 1. 只看到 1 张 GPU 在跑

✅ 没用 DistributedSampler

❌ 2. NCCL timeout

✅ 原因：

• GPU 被占用
• IB / 网络问题

解决：

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

❌ 3. 多机连不上

✅ 检查：

ping MASTER_ADDR
telnet MASTER_ADDR MASTER_PORT

八、推荐完整配置组合

PyTorch 2.x
CUDA 11.8 / 12.1
NCCL 2.18+
Ubuntu 22.04
torchrun + DDP

九、进阶（可选）

十、我可以进一步帮你

你可以直接告诉我：
1️⃣ 单机还是多机？
2️⃣ 显卡数量 & 型号？
3️⃣ 是否用 Docker / Slurm？
4️⃣ 训练 CV / NLP / LLM？
我可以给你 完全可运行的模板代码 + 启动脚本。