linux pytorch如何分布式

在 Linux 上使用 PyTorch 进行分布式训练，常见有 三种主流方式。下面按「从简单到专业」给你系统讲清楚，并附带可运行示例。

一、最常用：DistributedDataParallel（DDP，推荐✅）

适合：多卡 / 多机训练，性能最好

1️⃣ 核心概念

• torch.distributed
• 通信后端：nccl（GPU）、gloo（CPU）
• 每个进程 = 1 张 GPU

2️⃣ 单机多卡示例（最典型）

启动方式（推荐）

torchrun --nproc_per_node=4 train.py

train.py

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    for x, y in loader:
        x, y = x.cuda(local_rank), y.cuda(local_rank)
        loss = model(x, y)
        loss.backward()
        optimizer.step()

if __name__ == "__main__":
    main()

✅ 关键点

• 必须用 DistributedSampler
• batch_size 是 每张卡 的
• 不要手动reduce loss

二、简单并行：DataParallel（不推荐⚠️）

model = torch.nn.DataParallel(model.cuda())

❌ 问题：

• 只支持单机
• 主卡瓶颈
• 已逐渐淘汰

三、多机分布式（集群）

1️⃣ 环境假设

2️⃣ 启动命令

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

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

四、CPU / 跨节点 Gloo 后端

dist.init_process_group(backend="gloo")

适合：

• 无 GPU
• 调试

五、常见坑（非常重要⚠️）

❌ 1. 多进程 print 乱

if dist.get_rank() == 0:
    print(...)

❌ 2. 随机种子不一致

torch.manual_seed(42)

❌ 3. DataLoader num_workers

• 多进程下 num_workers 不要太大

六、如何判断是否在分布式环境？

dist.is_available() and dist.is_initialized()

七、总结选型

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

• ✅ 你是 单机几卡 / 多机
• ✅ 是否有 SLURM / Kubernetes
• ✅ 用的是 CNN / LLM / 微调

我可以直接给你 可拷贝的完整训练模板。