如何配置PyTorch分布式环境

配置 PyTorch 分布式环境主要包括 硬件/网络准备、软件环境、进程启动方式、代码编写 四个部分。下面按常见场景（单机多卡 / 多机多卡）系统说明。

一、环境准备

1. 硬件与网络

• 单机多卡：一台机器 + 多张 GPU（NVIDIA）
• 多机多卡：多台机器 + 高速网络（建议 InfiniBand / RoCE）
• 确保：
• 所有机器 GPU 驱动、CUDA 版本一致
• 机器之间 SSH 免密互通
• 防火墙未阻断 NCCL 通信端口（或统一规划端口）

2. 软件环境

Python & PyTorch

# 推荐 Python >= 3.8
conda create -n torch_dist python=3.10
conda activate torch_dist
# 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.distributed.is_available())

NCCL（通常 PyTorch 自带）

# 查看 NCCL 版本
python -c "import torch; print(torch.cuda.nccl.version())"

二、分布式模式选择

✅ 推荐：torch.distributed + NCCL

三、分布式初始化方式

1. 环境变量方式（最常用）

import os
import torch
import torch.distributed as dist
def setup():
dist.init_process_group(
backend="nccl",
init_method="env://"
)

需要设置环境变量：

export MASTER_ADDR="127.0.0.1"
export MASTER_PORT=29500
export WORLD_SIZE=4
export RANK=0

2. torchrun（强烈推荐 ✅）

不需要手动设置 RANK / WORLD_SIZE

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

四、单机多卡示例（DDP）

1. 代码结构

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

2. 示例 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
from torch.utils.data.distributed import 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(128, 10).cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
dataset = torch.randn(1000, 128)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(3):
sampler.set_epoch(epoch)
for x in loader:
x = x.cuda(local_rank)
loss = model(x).sum()
loss.backward()
optimizer.step()
optimizer.zero_grad()
cleanup()
if __name__ == "__main__":
main()

3. 启动

torchrun --nproc_per_node=4 train.py

五、多机多卡配置

假设：

• 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

六、常见坑 & 排查

1. 卡住（hang）

• init_process_group 没执行完
• WORLD_SIZE 不一致
• 防火墙未关闭

2. NCCL 报错

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

3. 数据加载

✅ 必须用 DistributedSampler

❌ 不要用 shuffle=True

七、进阶配置

• FSDP（超大模型）
• 混合精度（AMP）
• 梯度累积
• Checkpoint 保存（只保存 rank 0）

如果你愿意，我可以：

• ✅ 给你 FSDP 配置示例
• ✅ 帮你 迁移现有单卡代码到 DDP
• ✅ 针对 你当前 GPU / 集群环境定制配置

直接告诉我你的 GPU 数量 + 单机/多机 + PyTorch 版本 即可。