PyTorch在CentOS上的分布式训练指南

一环境准备与系统检查

驱动与工具链
- 确认已安装 NVIDIA 驱动 与 CUDA Toolkit，使用 nvidia-smi 检查 GPU 状态与驱动版本；确保 PyTorch 与 CUDA 版本匹配。
网络与账户
- 多机训练时，各节点需在同一二层/三层网络，建议使用 内网 IP；配置 SSH 免密登录（便于启动脚本分发与调试）。
防火墙与端口
- 开放用于进程组初始化的 MASTER_PORT（如 12345），或临时关闭防火墙以排除干扰：sudo systemctl stop firewalld（生产环境请按需放行端口）。
软件环境
- 建议使用 Python 3.8+、虚拟环境（conda/venv），各节点 PyTorch/CUDA 版本保持一致，避免因版本不一致导致通信或算子问题。

二方案一原生 PyTorch DDP 快速上手

关键概念
- 使用 DistributedDataParallel（DDP） 进行数据并行；GPU 训练推荐后端 NCCL；通过 init_process_group 建立进程组，常用 TCP 初始化方式或 env:// 环境变量方式。

最小可用示例（单文件可多进程启动）

保存为 ddp_demo.py：

import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset

def main(rank, world_size):
    # 1) 初始化进程组（TCP 方式）
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://<MASTER_IP>:<MASTER_PORT>',
        world_size=world_size,
        rank=rank
    )
    torch.cuda.set_device(rank)

    # 2) 模型与 DDP
    model = nn.Linear(10, 1).to(rank)
    model = DDP(model, device_ids=[rank])

    # 3) 数据（示例）
    ds = TensorDataset(torch.randn(1024, 10), torch.randn(1024, 1))
    sampler = DistributedSampler(ds, num_replicas=world_size, rank=rank)
    dl = DataLoader(ds, batch_size=32, sampler=sampler, num_workers=4, pin_memory=True)

    # 4) 训练
    opt = optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()
    for epoch in range(3):
        sampler.set_epoch(epoch)
        for x, y in dl:
            x, y = x.to(rank), y.to(rank)
            opt.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()
            opt.step()
    dist.destroy_process_group()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--rank', type=int, required=True)
    parser.add_argument('--world-size', type=int, required=True)
    args = parser.parse_args()
    main(args.rank, args.world_size)

启动方式

单机多卡（示例为 4 卡）：

python -m torch.distributed.launch \
  --nproc_per_node=4 \
  --nnodes=1 --node_rank=0 \
  --master_addr=<MASTER_IP> --master_port=12345 \
  ddp_demo.py --rank 0 --world-size 4

多机多卡（示例为 2 节点，每节点 4 卡）：

节点0（rank 0）：

python -m torch.distributed.launch \
  --nproc_per_node=4 \
  --nnodes=2 --node_rank=0 \
  --master_addr=<MASTER_IP> --master_port=12345 \
  ddp_demo.py --rank 0 --world-size 8

节点1（rank 1）：

python -m torch.distributed.launch \
  --nproc_per_node=4 \
  --nnodes=2 --node_rank=1 \
  --master_addr=<MASTER_IP> --master_port=12345 \
  ddp_demo.py --rank 1 --world-size 8

提示：将 <MASTER_IP> 设为主节点的内网地址，确保各节点可互通；必要时用环境变量方式传参（init_method=‘env://’）。

三方案二 Horovod 多机多卡

适用场景与特点
- Horovod 基于 Ring-AllReduce，对 PyTorch/TensorFlow/Keras 提供统一接口；安装便捷，支持 MPI/Gloo 通信后端，易于横向扩展。

安装与启动

安装（示例）：pip install horovod[pytorch]（需提前装好 OpenMPI 或 Gloo）。

训练启动（示例，使用 mpirun 与 OpenMPI）：

mpirun -np 8 -H node0:4,node1:4 \
  -bind-to none -map-by slot \
  -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH \
  python train_horovod.py

代码改造要点（PyTorch）

import torch
import horovod.torch as hvd
hvd.init()
torch.cuda.set_device(hvd.local_rank())

model = Net().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01 * hvd.size())
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

# 广播初始参数
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

# 数据：每个进程仅读取自己 shard
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
)
train_loader = DataLoader(train_dataset, batch_size=..., sampler=train_sampler, ...)

说明：Horovod 适合已有单卡脚本快速扩展到多机多卡，通信与容错机制成熟。

四性能与稳定性要点

数据与采样
- 使用 DistributedSampler，并在每个 epoch 调用 sampler.set_epoch(epoch)，保证各进程数据划分随机且可复现。
批大小与学习率
- 多卡训练时全局 batch_size = 每卡 batch_size × GPU 总数；学习率通常按 GPU 数 线性或按 sqrt(卡数) 缩放（视优化器与任务而定）。
数据加载
- 设置 num_workers、开启 pin_memory=True，加速 CPU→GPU 传输；确保 共享存储（如 NFS/并行文件系统）带宽与 I/O 不成为瓶颈。
通信与后端
- GPU 训练优先 NCCL；如跨节点网络一般或需 CPU 训练，可改用 Gloo。
稳定性
- 固定随机种子（Python/torch/cuda），统一 NCCL/驱动/CUDA/PyTorch 版本；必要时开启 NCCL_DEBUG=INFO 定位通信问题。

五常见问题排查

进程卡在 init_process_group
- 检查 MASTER_ADDR/MASTER_PORT 是否一致且可达；确认 防火墙/安全组 已放行端口；各节点 world_size/rank 配置无误。
NCCL 超时/错误
- 设置环境变量：export NCCL_SOCKET_TIMEOUT=3600、export NCCL_DEBUG=INFO；检查 驱动/CUDA/NCCL 版本匹配与网络健康度。
SSH 免密失败
- 确保 ~/.ssh/authorized_keys 正确、权限为 600/700；必要时用 ssh-copy-id 分发公钥。
版本不一致导致异常
- 各节点统一 PyTorch/CUDA/cuDNN 版本；容器化训练可显著降低环境差异风险。
单机可用多机失败
- 多机需使用 内网 IP，避免 127.0.0.1；确认路由与交换机策略允许节点间通信。

最新问答

相关标签