Linux环境下PyTorch集群搭建指南

一、前置准备

1. 硬件要求

至少2台Linux服务器（推荐CentOS/Ubuntu），其中1台作为主节点（Master），其余作为工作节点（Worker）；
所有节点配置相同操作系统版本（避免兼容性问题）；
建议每台节点配备GPU（若需GPU加速训练），且GPU型号一致（如NVIDIA A100）。

2. 软件与环境要求

所有节点安装相同版本的Python（推荐3.8+）和pip；

安装系统依赖（编译工具、开发库）：

# CentOS系统
sudo yum update -y
sudo yum install -y gcc-c++ make cmake git
# Ubuntu系统
sudo apt-get update && sudo apt-get install -y build-essential cmake git

3. 网络配置

所有节点处于同一局域网，确保IP可达（通过ping <节点IP>测试）；

关闭防火墙或开放通信端口（如23456，用于节点间同步）：

# CentOS系统
sudo firewall-cmd --zone=public --add-port=23456/tcp --permanent
sudo firewall-cmd --reload
# Ubuntu系统（ufw）
sudo ufw allow 23456/tcp

4. SSH无密码登录

在主节点生成SSH密钥对：

ssh-keygen -t rsa  # 直接回车，默认保存路径~/.ssh/id_rsa

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

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

验证无密码登录：

ssh user@worker1_ip  # 无需输入密码即可登录

二、PyTorch安装

1. 安装PyTorch

CPU版本（适用于无GPU场景）：

pip3 install torch torchvision torchaudio

GPU版本（需提前安装CUDA/cuDNN）：

根据CUDA版本选择对应命令（如CUDA 11.3）：

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

验证安装：

import torch
print(torch.cuda.is_available())  # 应输出True

2. 可选工具安装

分布式任务调度：安装Slurm（适合大规模集群）或Dask（轻量级分布式计算）；
```
pip3 install dask distributed  # Dask安装
```
MPI支持（可选，用于更高效的进程通信）：
```
pip3 install mpi4py
```

三、分布式训练配置

1. 初始化进程组

在训练脚本中，使用torch.distributed.init_process_group初始化分布式环境：

import torch.distributed as dist

def setup(rank, world_size):
    # 初始化进程组，推荐使用NCCL后端（GPU加速）
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://<master_ip>:<master_port>',  # 主节点IP和端口
        world_size=world_size,  # 总进程数（节点数×每个节点GPU数）
        rank=rank  # 当前进程的全局排名（0到world_size-1）
    )

def cleanup():
    dist.destroy_process_group()  # 训练结束后销毁进程组

2. 数据并行封装

使用torch.nn.parallel.DistributedDataParallel（DDP）包装模型，实现数据并行：

import torch.nn as nn

model = YourModel().to(rank)  # 将模型移动到当前GPU
ddp_model = nn.parallel.DistributedDataParallel(
    model,
    device_ids=[rank]  # 当前节点的GPU索引
)

3. 数据加载调整

使用DistributedSampler确保每个进程加载不同的数据子集：

from torch.utils.data import DataLoader, DistributedSampler

dataset = YourDataset()  # 自定义数据集
sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,  # 总进程数
    rank=rank  # 当前进程排名
)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler  # 使用分布式采样器
)

4. 训练循环修改

在训练循环中，每轮迭代前调用sampler.set_epoch(epoch)，确保数据打乱顺序：

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # 每轮重置采样器，避免数据重复
    for data, target in dataloader:
        data, target = data.to(rank), target.to(rank)
        optimizer.zero_grad()
        output = ddp_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

四、启动分布式训练

1. 使用torch.distributed.launch工具

在主节点运行以下命令，启动分布式训练：

python -m torch.distributed.launch \
    --nproc_per_node=<num_gpus> \  # 每个节点的GPU数量（如4）
    --nnodes=<total_nodes> \       # 总节点数（如2）
    --node_rank=<current_node_rank> \  # 当前节点排名（主节点0，工作节点1、2...）
    --master_addr=<master_ip> \    # 主节点IP
    --master_port=<port> \         # 主节点端口（如23456）
    your_training_script.py        # 训练脚本路径

示例（2节点，每节点4个GPU）：

主节点命令：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.100" --master_port=23456 train.py

工作节点命令：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.168.1.100" --master_port=23456 train.py

2. 使用启动脚本（简化操作）

编写启动脚本start_train.sh，自动分发命令到各节点：

#!/bin/bash
# 主节点IP和端口
MASTER_ADDR="192.168.1.100"
MASTER_PORT=23456
# 总节点数
NNODES=2
# 每个节点的GPU数量
GPUS_PER_NODE=4

# 主节点运行
if [ "$1" == "master" ]; then
    echo "Starting master node..."
    python -m torch.distributed.launch \
        --nproc_per_node=$GPUS_PER_NODE \
        --nnodes=$NNODES \
        --node_rank=0 \
        --master_addr=$MASTER_ADDR \
        --master_port=$MASTER_PORT \
        train.py
# 工作节点运行
elif [ "$1" == "worker" ]; then
    echo "Starting worker node..."
    python -m torch.distributed.launch \
        --nproc_per_node=$GPUS_PER_NODE \
        --nnodes=$NNODES \
        --node_rank=1 \
        --master_addr=$MASTER_ADDR \
        --master_port=$MASTER_PORT \
        train.py
else
    echo "Usage: $0 {master|worker}"
fi

给脚本添加执行权限：
```
chmod +x start_train.sh
```
主节点运行：
```
./start_train.sh master
```
工作节点运行：
```
./start_train.sh worker
```

五、验证与监控

1. 验证集群功能

在训练脚本中添加日志，打印每个进程的rank和world_size：
```
print(f"Rank {rank}, World Size {world_size}")
```
运行后，检查所有节点是否输出正确的rank（0到world_size-1）。

2. 监控工具

Dask Dashboard：若使用Dask，通过http://<master_ip>:8787查看任务进度和集群状态；

PyTorch Profiler：使用torch.profiler监控训练性能，识别瓶颈：

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    on_trace_ready=lambda prof: prof.export_chrome_trace("trace.json")
) as prof:
    for data, target in dataloader:
        # 训练代码
        prof.step()

六、注意事项

1. 时间同步

所有节点必须同步时间（使用NTP服务），否则可能导致进程同步失败：

# CentOS系统
sudo yum install -y ntp
sudo systemctl start ntpd
sudo systemctl enable ntpd
# Ubuntu系统
sudo apt-get install -y ntp
sudo systemctl start ntp
sudo systemctl enable ntp

2. 环境一致性

使用requirements.txt或environment.yml统一所有节点的Python环境：

pip freeze > requirements.txt  # 导出主节点环境
# 在工作节点安装相同环境
pip install -r requirements.txt

3. 资源分配

根据节点GPU数量调整--nproc_per_node，避免资源浪费；
若使用多机多卡，确保world_size等于“节点数×每个节点GPU数”。

4. 故障排查

检查节点间通信：ping <节点IP>；
检查SSH登录：ssh user@worker_ip；
查看日志：训练脚本的输出日志（如train.log），定位错误信息。

Linux PyTorch集群搭建指南