PyTorch在CentOS上的分布式训练实践

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

一、环境准备

1. 系统与依赖配置

   • 安装基础工具：sudo yum install -y epel-release（启用EPEL仓库），sudo yum install -y libnccl-devel（NCCL后端依赖，用于多GPU通信）。
   • 安装Python环境管理工具：推荐使用Miniconda（避免系统Python冲突），通过wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh下载并安装。
   • 安装PyTorch：根据CUDA版本选择安装命令（如CUDA 11.8）：conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia。

2. 网络与SSH配置

   • 静态IP与端口开放：确保所有节点（主节点+工作节点）配置静态IP，防火墙开放训练所需的端口（如MASTER_PORT=10086），命令：sudo firewall-cmd --add-port=10086/tcp --permanent → sudo firewall-cmd --reload。
   • SSH无密码登录：在主节点生成密钥（ssh-keygen -t rsa），将公钥复制到所有工作节点（ssh-copy-id user@worker_ip），测试登录（ssh user@worker_ip）是否无需密码。

二、编写分布式训练脚本

以下是一个基于DistributedDataParallel（DDP）的模板脚本（train_ddp.py）：

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
import torch.distributed as dist

def setup(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = '主节点IP'  # 如'192.168.1.100'
    os.environ['MASTER_PORT'] = '10086'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)  # 将当前进程绑定到指定GPU

def cleanup():
    """清理分布式环境"""
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    """示例模型（替换为实际模型）"""
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)
    def forward(self, x):
        return self.fc(x.view(-1, 784))

def train(rank, world_size, epochs=5):
    setup(rank, world_size)
    
    # 数据加载（使用DistributedSampler确保数据分片）
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)
    
    # 模型与优化器（DDP包装模型）
    model = SimpleModel().to(rank)
    ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])
    criterion = nn.CrossEntropyLoss().to(rank)
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
    
    # 训练循环
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # 每个epoch打乱数据分片
        running_loss = 0.0
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(rank), labels.to(rank)
            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        if rank == 0:  # 仅主节点打印日志
            print(f"Epoch {
epoch+1}
, Loss: {
running_loss/len(dataloader)}
    ")
    
    cleanup()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--world_size", type=int, default=2, help="总进程数（节点数×每节点GPU数）")
    parser.add_argument("--rank", type=int, default=0, help="当前进程的全局排名")
    args = parser.parse_args()
    train(args.rank, args.world_size)

关键点说明：

• setup()与cleanup()：封装分布式环境初始化与清理，避免代码重复。
• DistributedSampler：确保每个GPU处理不同的数据分片，避免数据重复。
• DistributedDataParallel（DDP）：自动处理梯度同步与模型更新，是PyTorch推荐的分布式训练方式。

三、启动分布式训练

1. 单机多卡（1台机器，多块GPU）

使用torch.distributed.launch工具（PyTorch内置），命令格式：

python -m torch.distributed.launch \
    --nproc_per_node=NUM_GPUS_PER_NODE \  # 每个节点的GPU数量（如2）
    --nnodes=1 \                         # 节点数量（单机为1）
    --node_rank=0 \                      # 当前节点排名（单机主节点为0）
    --master_addr="主节点IP" \           # 主节点IP（单机为自己）
    --master_port=MASTER_PORT \          # 主节点端口（如10086）
    train_ddp.py

示例（单机2张GPU）：

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 --node_rank=0 --master_addr="192.168.1.100" --master_port=10086 train_ddp.py

2. 多机多卡（多台机器，多块GPU）

在每台节点上运行相同的命令，通过--node_rank区分节点身份（主节点为0，工作节点依次递增）。
主节点（node0）：

python -m torch.distributed.launch \
    --nproc_per_node=4 \  # 主节点有4张GPU
    --nnodes=2 \          # 总节点数（主节点+1个工作节点）
    --node_rank=0 \       # 主节点排名
    --master_addr="192.168.1.100" \  # 主节点IP
    --master_port=10086 \  # 主节点端口
    train_ddp.py

工作节点（node1）：

python -m torch.distributed.launch \
    --nproc_per_node=2 \  # 工作节点有2张GPU
    --nnodes=2 \          # 总节点数
    --node_rank=1 \       # 工作节点排名
    --master_addr="192.168.1.100" \  # 主节点IP（与主节点一致）
    --master_port=10086 \  # 主节点端口（与主节点一致）
    train_ddp.py

注意：--nproc_per_node必须与节点实际的GPU数量一致，否则会报错。

四、注意事项

1. 版本一致性：所有节点的PyTorch、CUDA、cuDNN版本必须完全一致，避免兼容性问题。
2. 资源监控：使用nvidia-smi实时查看GPU利用率（如watch -n 1 nvidia-smi），确保GPU资源被充分利用。
3. 日志管理：仅在主节点（rank=0）保存模型（if rank == 0: torch.save(model.state_dict(), "model.pth")），避免重复保存。
4. 调试技巧：先在单机多卡环境下测试脚本逻辑（如数据加载、模型前向传播），确认无误后再扩展到多机多卡。

通过以上步骤，即可在CentOS系统上实现PyTorch的分布式训练，充分利用多机多卡的算力加速模型训练过程。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！