PyTorch分布式训练如何配置

下面从概念 → 常用方式 → 配置步骤 → 示例四个层面，系统讲一下 PyTorch 分布式训练如何配置。

一、PyTorch 分布式核心概念

PyTorch 分布式训练主要依赖：

• Process Group（进程组）：一组参与训练的进程
• Backend（后端）：
• nccl（GPU，推荐）
• gloo（CPU / 部分 GPU）
• mpi（需要 MPI 环境）
• 通信方式：
• DP（DataParallel）：单机多卡，已不推荐
• DDP（DistributedDataParallel）：主流方案 ✅
• FSDP（Fully Sharded Data Parallel）：大模型训练

二、常见分布式训练模式

1️⃣ 单机多卡（最常用）

1 台机器 × 多张 GPU

• 使用 torch.distributed.launch 或 torchrun
• backend = nccl

2️⃣ 多机多卡

多台机器 × 多张 GPU

• 需要：
• 固定 MASTER_ADDR
• 固定 MASTER_PORT
• 各节点 rank 不同

三、DDP 标准配置流程（重点）

✅ 1. 初始化进程组

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

✅ 2. 设置本地 rank（关键）

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

✅ 3. 模型封装为 DDP

model = MyModel().cuda()
model = torch.nn.parallel.DDistributedDataParallel(
model,
device_ids=[local_rank]
)

✅ 4. 使用 DistributedSampler

from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset)
dataloader = DataLoader(
dataset,
batch_size=32,
sampler=sampler
)

⚠️ 不要使用 shuffle=True

四、启动方式（非常重要）

✅ 推荐方式：torchrun（PyTorch ≥ 1.9）

单机 4 卡

torchrun \
--nproc_per_node=4 \
train.py

多机（2 机 × 4 卡）

# 节点 0
torchrun \
--nnodes=2 \
--node_rank=0 \
--nproc_per_node=4 \
--master_addr=192.168.1.1 \
--master_port=29500 \
train.py
# 节点 1
torchrun \
--nnodes=2 \
--node_rank=1 \
--nproc_per_node=4 \
--master_addr=192.168.1.1 \
--master_port=29500 \
train.py

五、完整 DDP 示例（简化版）

import os
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
def main():
local_rank = int(os.environ["LOCAL_RANK"])
dist.init_process_group("nccl")
torch.cuda.set_device(local_rank)
dataset = MyDataset()
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
model = MyModel().cuda()
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank]
)
for x, y in loader:
x, y = x.cuda(), y.cuda()
loss = model(x, y)
loss.backward()
optimizer.step()
if __name__ == "__main__":
main()

六、常见坑 & 建议

✅ 必须使用 DistributedSampler
✅ batch_size 是 单卡 batch
✅ 保存模型只在 rank == 0
✅ 日志只在 rank == 0 打印

✅ 多机训练确保网络互通（端口可访问）

七、什么时候用 FSDP？

• 模型大到单卡放不下
• 参数量 > 单卡显存
• 使用 torch.distributed.fsdp

如果你愿意，我可以：

• 给你 多机 GPU 集群配置模板
• 讲 DDP vs FSDP 对比
• 帮你 debug 分布式训练报错

你现在的训练场景是？