如何在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數量及批量大小，以獲得最佳性能。