PyTorch分布式训练如何扩容

在 PyTorch 中进行分布式训练“扩容”（scale out / scale up），本质是把更多计算资源（GPU / 节点）高效利用起来，同时尽量少改代码、少踩坑。下面按思路 → 方法 → 实操建议来讲。

一、先明确：你要扩的是什么？

二、PyTorch 主流分布式方式（选型）

1️⃣ DataParallel（不推荐扩容）

model = nn.DataParallel(model)

❌ 缺点：

• 仅支持单机
• GIL + Python 主线程瓶颈
• 通信效率低

✅ 只适合：debug / 小实验

2️⃣ DistributedDataParallel（✅ 主流方案）

torch.nn.parallel.DistributedDataParallel

✅ 支持：

• 单机多卡
• 多机多卡
• 高吞吐、低通信开销

✅ 适合 90% 的扩容场景

3️⃣ FSDP / FullyShardedDataParallel（大模型必备）

torch.distributed.fsdp.FullyShardedDataParallel

✅ 特点：

• 参数 + 梯度 + 优化器状态分片
• 能训 远超单卡显存 的模型

✅ 适合：

• LLM
• 多卡 / 多机大模型训练

三、扩容的核心步骤（以 DDP 为例）

1️⃣ 初始化进程组

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

2️⃣ 启动方式（关键）

单机多卡

torchrun --nproc_per_node=8 train.py

多机多卡

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

3️⃣ 模型 & 数据改造

模型

model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

数据（非常重要）

sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler)

⚠️ 不加重 sampler → 数据重复 → 白扩

四、扩容时的关键调优点

✅ 1. Batch Size 线性放大

• 8 GPU → batch ×8
• 学习率通常 ×√N 或 ×N（看任务）

✅ 2. 通信优化

• backend="nccl"
• 使用 InfiniBand / RDMA
• gradient_as_bucket_view=True
• static_graph=True（可选）

✅ 3. 梯度累积（显存不够时）

loss = loss / accum_steps
loss.backward()
if step % accum_steps == 0:
optimizer.step()

✅ 4. 日志 & Checkpoint

• 只让 rank == 0 写
• 用 torch.save(model.module.state_dict())

五、从 1 机 → N 机的扩容清单

✅ 网络互通（SSH / 防火墙）
✅ NCCL 环境变量

export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_DISABLE=0

✅ 共享存储（NFS / S3 / HDFS）

✅ 固定 master_addr / port

六、什么时候该用 FSDP？

七、一句话总结

PyTorch 扩容 = DDP/FSDP + torchrun + DistributedSampler + 网络与显存优化
如果你愿意，可以告诉我：

• 你是 单机多卡 / 多机
• 模型大小
• 是否训练 LLM

我可以直接给你可运行的模板代码。