PyTorch在CentOS上的分布式训练实践指南
系统与依赖配置
sudo yum install -y epel-release（启用EPEL仓库），sudo yum install -y libnccl-devel（NCCL后端依赖，用于多GPU通信）。wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh下载并安装。conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia。网络与SSH配置
MASTER_PORT=10086），命令：sudo firewall-cmd --add-port=10086/tcp --permanent → sudo firewall-cmd --reload。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推荐的分布式训练方式。使用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 在每台节点上运行相同的命令，通过--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数量一致，否则会报错。
nvidia-smi实时查看GPU利用率（如watch -n 1 nvidia-smi），确保GPU资源被充分利用。rank=0）保存模型（if rank == 0: torch.save(model.state_dict(), "model.pth")），避免重复保存。通过以上步骤，即可在CentOS系统上实现PyTorch的分布式训练，充分利用多机多卡的算力加速模型训练过程。
