PyTorch分布式训练怎么部署

下面系统性地讲一下 PyTorch 分布式训练的部署方式，从概念到常见落地方案，偏实战导向。

一、PyTorch 分布式训练的核心概念

PyTorch 分布式训练主要依赖 torch.distributed，常见模式有两种：

1️⃣ DP（DataParallel）【不推荐】

• 单进程、多 GPU
• 主 GPU 易成为瓶颈
• 仅适合小规模调试

model = torch.nn.DataParallel(model)

2️⃣ DDP（DistributedDataParallel）【推荐】

• 多进程、多 GPU
• 每个 GPU 一个进程
• 效率高、可扩展性强

model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank]
)

二、分布式训练的关键组件

1️⃣ 通信后端（Backend）

torch.distributed.init_process_group(backend="nccl")

2️⃣ 进程启动方式

• torchrun（推荐）
• multiprocessing
• Slurm（集群）

三、单机多卡部署（最常用）

示例：单机 8 卡

torchrun \
--nproc_per_node=8 \
train.py

train.py 关键代码：

import torch
import torch.distributed as dist
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = model.cuda(local_rank)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank]
)

四、多机多卡部署（集群）

1️⃣ 环境假设

• 2 台机器
• 每台 8 张 GPU
• 机器 IP：
• node0: 192.168.1.10
• node1: 192.168.1.11

2️⃣ 启动命令

Node 0

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

Node 1

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

五、数据并行要点

✅ 使用 DistributedSampler

from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler)

⚠️ 不要使用 shuffle=True

六、容器化部署（生产推荐）

Docker + PyTorch 分布式

FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
COPY . /app
WORKDIR /app
CMD ["torchrun", "train.py"]

Kubernetes（高级）

• 使用 PyTorchJob（Kubeflow）
• 自动管理 master / worker
• 适合大规模训练

七、常见部署方案对比

八、新手常见坑

❌ 忘记 set_device
❌ 没用 DistributedSampler
❌ master_port 被占用
❌ 多机网络不通（防火墙）

❌ 日志只在一个进程打印

九、总结一句话

PyTorch 分布式训练 = DDP + torchrun + DistributedSampler
如果你愿意，可以告诉我：

• ✅ 单机 / 多机
• ✅ GPU 数量
• ✅ 是否用 Docker / Slurm / K8s

我可以直接给你一套可复制的部署模板。