如何在Ubuntu上搭建PyTorch集群

在Ubuntu上搭建PyTorch集群的详细步骤

1. 硬件与环境准备

服务器配置：至少需要2台Ubuntu系统服务器（建议配置相同，如Ubuntu 22.04 LTS），每台服务器配备1块及以上NVIDIA GPU（需支持CUDA）。

基础软件包：在所有节点上安装Python 3.8+、pip及系统依赖：

sudo apt update && sudo apt install -y python3 python3-pip python3-dev build-essential

PyTorch安装：通过PyTorch官网获取对应CUDA版本的安装命令（如CUDA 11.8）：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

依赖工具：安装mpi4py（用于MPI通信，可选但推荐）：
```
pip3 install mpi4py
```

2. 网络与SSH配置

网络连通性：确保所有节点之间能互相ping通（如ping node2_ip），关闭防火墙或开放PyTorch使用的端口（默认12345）：
```
sudo ufw allow 12345/tcp  # 若使用ufw防火墙
```
SSH无密码登录：在主节点生成SSH密钥，并分发到所有工作节点：
```
ssh-keygen -t rsa  # 主节点执行，按回车生成密钥
ssh-copy-id user@node2_ip  # 将公钥复制到工作节点（替换user和node2_ip）
```
测试无密码登录：ssh user@node2_ip，无需输入密码即为成功。

3. 集群参数配置

环境变量设置：在所有节点上配置分布式训练所需的环境变量（以4个GPU为例）：

export MASTER_ADDR='主节点IP'  # 如192.168.1.10
export MASTER_PORT=12345       # 未被占用的端口
export WORLD_SIZE=4            # 总GPU数量（主节点+工作节点的GPU数之和）
export NCCL_DEBUG=INFO         # 开启NCCL调试信息（可选）
export NCCL_SOCKET_IFNAME=eth0 # 指定通信网卡（根据实际情况修改，如eth0、ens33）

可将上述命令添加到~/.bashrc中，避免每次重启终端重复设置。

4. 编写分布式训练脚本

创建distributed_train.py，核心逻辑包括进程组初始化、DDP模型包装、分布式数据加载：

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
import torchvision.datasets as datasets
import torchvision.transforms as transforms

def main():
    # 初始化进程组（使用环境变量自动获取参数）
    torch.distributed.init_process_group(backend='nccl', init_method='env://')
    rank = torch.distributed.get_rank()  # 当前进程的rank（0~WORLD_SIZE-1）
    torch.cuda.set_device(rank)          # 设置当前进程使用的GPU

    # 定义模型并移动到对应GPU
    model = nn.Sequential(
        nn.Linear(784, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    ).to(rank)
    model = DDP(model, device_ids=[rank])  # 包装模型以实现梯度同步

    # 定义损失函数与优化器
    criterion = nn.CrossEntropyLoss().to(rank)
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # 加载数据集（使用DistributedSampler分割数据）
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=torch.distributed.get_world_size(), rank=rank)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # 训练循环
    for epoch in range(5):
        sampler.set_epoch(epoch)  # 每个epoch打乱数据分布
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(loader):
            inputs, labels = inputs.to(rank), labels.to(rank)
            optimizer.zero_grad()
            outputs = model(inputs.view(-1, 784))
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        if rank == 0:  # 仅主节点打印日志
            print(f'Epoch {epoch+1}, Loss: {running_loss/len(loader)}')

    # 清理进程组
    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    main()

5. 启动分布式训练

单节点多GPU训练（如4个GPU）：
```
python -m torch.distributed.launch --nproc_per_node=4 distributed_train.py
```
--nproc_per_node指定每个节点使用的GPU数量。

多节点训练（如主节点+2个工作节点，共4个GPU）：

# 主节点执行（rank=0）
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="主节点IP" --master_port=12345 distributed_train.py

# 工作节点1执行（rank=1）
python -m torch.distributed.launch --nproc_per_node=1 --nnodes=3 --node_rank=1 --master_addr="主节点IP" --master_port=12345 distributed_train.py

# 工作节点2执行（rank=2）
python -m torch.distributed.launch --nproc_per_node=1 --nnodes=3 --node_rank=2 --master_addr="主节点IP" --master_port=12345 distributed_train.py

参数说明：

--nnodes：集群中的节点总数（主节点+工作节点）。
--node_rank：当前节点的rank（0~nnodes-1，主节点为0）。
--master_addr：主节点的IP地址。
--master_port：主节点监听的端口。

6. 可选优化与监控

使用Accelerate简化配置：安装Hugging Face的accelerate库，自动处理环境变量与进程启动：

pip install accelerate
accelerate config  # 按提示配置集群参数
accelerate launch distributed_train.py  # 一键启动

监控工具：通过nccl-tests测试GPU通信性能，使用htop监控CPU/GPU利用率：

git clone https://github.com/NVIDIA/nccl-tests.git
cd nccl-tests && make MPI=1
mpirun -np 4 ./build/all_reduce_perf -b 8G -e 8G -f 2 -g 4

通过以上步骤，即可在Ubuntu上搭建PyTorch集群，实现高效的分布式训练。需注意根据实际环境调整网络配置、GPU数量及批量大小，以获得最佳性能。