PyTorch分布式训练怎么扩容

在 PyTorch 里“分布式训练扩容”通常指从单机多卡 → 多机多卡，或增加节点/GPU 数量后仍能高效训练。下面按「扩容方式 + 关键配置 + 常见坑」来说明。

一、PyTorch 分布式主流方案

PyTorch 分布式训练主要用：

1. DDP（DistributedDataParallel）（最常用）
2. FSDP（FullyShardedDataParallel）（大模型扩容必备）
3. DeepSpeed / Megatron（超大规模）

扩容本质：增加 rank 数量（GPU 数）

二、单机多卡 → 多机多卡（标准扩容）

1. 初始化方式（关键）

import torch.distributed as dist
dist.init_process_group(
backend="nccl",
init_method="tcp://MASTER_IP:29500",
rank=rank,
world_size=world_size
)

• MASTER_IP：主节点 IP
• world_size = 机器数 × 每机GPU数
• rank 全局唯一

2. 启动方式

方式一：torchrun（推荐）

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

✅ 扩容只改 --nnodes 和 --node_rank

三、扩容时数据怎么处理

DataLoader 必须改

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

⚠️ 注意：

• batch_size 是每卡 batch
• 总 batch = batch_size × world_size

四、模型并行扩容（显存不够）

1. DDP（数据并行）

• 每卡一份模型
• 适合中等模型

2. FSDP（参数分片）

from torch.distributed.fsdp import FSDP
model = FSDP(model)

✅ 扩容 GPU 数 → 显存线性下降

✅ 适合 LLM / 大模型

五、通信与性能扩容要点

1. 网络

• 多机必须：RDMA / InfiniBand
• 否则通信成为瓶颈

2. 梯度通信

• DDP 自动 AllReduce
• 可开启：

export NCCL_ASYNC_ERROR_HANDLING=1

六、弹性扩容（动态加节点）

方案

• torchelastic

torchrun --nnodes=1:4 --nproc_per_node=8 train.py

支持 1~4 节点动态加入
⚠️ 需要：

• 检查点频繁保存
• 支持 restart 逻辑

七、常见扩容坑

八、简单扩容经验公式

• LR ≈ 原LR × (新batch / 原batch)
• world_size 翻倍 → 吞吐接近翻倍（理想）

如果你愿意，可以告诉我：

• 你现在是 单机几卡
• 模型大小 / 是否 LLM
• 用 DDP 还是 FSDP

我可以给你一套直接可跑的扩容模板。