PyTorch CentOS集群怎么搭建

PyTorch CentOS集群搭建指南

1. 前提条件

• 至少2台CentOS服务器（1台主节点、1台及以上工作节点），操作系统版本一致；
• 所有节点之间可通过SSH无密码登录；
• 基础Linux操作知识（如vim、scp、systemctl）。

2. 环境准备

2.1 系统更新与依赖安装

在所有节点执行以下命令，更新系统并安装必要工具：

sudo yum update -y                  # 更新系统
sudo yum install -y python3 python3-pip gcc-c++ make cmake git  # 安装Python、pip及编译工具

若需GPU支持，还需安装CUDA Toolkit（需匹配PyTorch版本）和NCCL库（用于多卡通信）：

sudo yum install -y epel-release    # 启用EPEL仓库
sudo yum install -y libnccl-devel   # 安装NCCL库（多卡通信加速）

通过pip安装PyTorch（CPU版或GPU版）：

pip3 install torch torchvision torchaudio          # CPU版
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118  # GPU版（以CUDA 11.8为例）

若使用MPI通信（可选），安装mpi4py：

pip3 install mpi4py

2.2 SSH免密登录配置

在主节点生成SSH密钥对（若已有可跳过）：

ssh-keygen -t rsa                   # 按回车键默认保存路径

将公钥复制到所有工作节点：

ssh-copy-id user@worker_node_ip     # 替换为工作节点IP和用户名

验证免密登录：

ssh user@worker_node_ip             # 应无需输入密码

设置SSH目录权限（所有节点）：

chmod 700 ~/.ssh                    # 限制SSH目录权限
chmod 600 ~/.ssh/authorized_keys    # 限制authorized_keys文件权限

3. 集群配置

3.1 节点信息收集

• 主节点IP（如192.168.1.100）、开放端口（如29400，需防火墙放行）；
• 工作节点IP列表（如192.168.1.101、192.168.1.102）；
• 每个节点的GPU数量（如主节点2张、工作节点1张）。

3.2 启动分布式训练

使用torchrun（推荐，替代torch.distributed.launch）启动训练脚本，示例如下：
主节点命令（node_rank=0）：

torchrun --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="192.168.1.100" --master_port=29400 train.py

工作节点命令（node_rank=1）：

torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="192.168.1.100" --master_port=29400 train.py

参数说明：

• --nproc_per_node：每台节点的进程数（通常等于GPU数量）；
• --nnodes：集群总节点数；
• --node_rank：当前节点的全局排名（主节点0，工作节点依次递增）；
• --master_addr：主节点IP；
• --master_port：主节点通信端口（需防火墙放行）。

4. 编写分布式训练脚本

训练脚本需使用torch.distributed.init_process_group初始化分布式环境，并用DistributedDataParallel（DDP）包装模型。示例如下：

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

def main():
    # 初始化分布式环境（通过环境变量读取配置）
    torch.distributed.init_process_group(backend="nccl", init_method="env://")
    
    # 获取当前进程信息
    rank = int(os.environ["RANK"])          # 当前进程全局排名
    local_rank = int(os.environ["LOCAL_RANK"])  # 本机进程排名（对应GPU ID）
    world_size = int(os.environ["WORLD_SIZE"])  # 总进程数
    
    # 设置GPU设备
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    
    # 打印节点信息
    print(f"[Rank {
rank}
/{
world_size}
] Using GPU {
local_rank}
: {
torch.cuda.get_device_name(local_rank)}
")
    
    # 构建模型并迁移至GPU
    model = nn.Linear(10, 10).to(device)
    ddp_model = DDP(model, device_ids=[local_rank])  # 包装模型（需指定本机GPU）
    
    # 构造数据集与数据加载器（使用DistributedSampler保证数据不重复）
    dataset = torch.utils.data.TensorDataset(torch.randn(100, 10), torch.randint(0, 10, (100,)))
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=10, sampler=sampler)
    
    # 定义损失函数与优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
    
    # 训练循环
    for epoch in range(5):
        sampler.set_epoch(epoch)  # 每个epoch重新打乱数据
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            # 前向传播
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            
            # 反向传播与优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        # 仅主节点打印损失
        if rank == 0:
            print(f"Epoch {
epoch}
, Loss: {
loss.item()}
    ")

if __name__ == "__main__":
    main()

关键点说明：

• init_process_group：通过env://自动读取环境变量（RANK、WORLD_SIZE等），无需手动设置；
• DistributedSampler：确保每个进程处理不同的数据子集，避免数据重复；
• DDP包装模型：自动处理梯度同步（All-Reduce操作），提升多卡训练效率。

5. 测试与验证

• 单节点多卡测试：在主节点运行torchrun --nproc_per_node=2 --nnodes=1 train.py，验证多GPU是否正常工作；
• 多节点测试：在主节点和工作节点分别运行上述命令，检查是否能同步训练（主节点应打印所有节点的损失）；
• 日志查看：通过tail -f log.txt（若脚本重定向日志）查看各节点输出，确认无通信错误。

6. 注意事项

• 环境一致性：所有节点的PyTorch、CUDA、NCCL版本必须完全一致（建议使用Docker容器化部署，避免环境冲突）；
• 网络配置：关闭节点间的防火墙或放行master_port（如sudo firewall-cmd --add-port=29400/tcp --permanent）；
• 资源分配：根据节点GPU数量调整--nproc_per_node，避免资源争用；
• 调试技巧：先在单节点多卡测试通过，再扩展至多节点，减少排查难度。

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