从单卡到集群的“加速”之痛

我刚刚接手了一个跨国电商推荐模型的训练任务。训练数据量接近1TB，单卡训练一个epoch需要整整16小时。业务侧每周至少要迭代3次模型效果，这意味着训练时间要从“等一夜”缩到“等一杯咖啡”的时间级别。经过对比测试，A5数据决定在CentOS 8服务器集群上部署PyTorch分布式训练框架。

CentOS 8虽然在2021年底后社区维护进入生命周期尾声，但在企业级机房仍然大量使用。我们要在其上实现高效的PyTorch分布式训练，涉及驱动、库依赖、网络、GPU和调度多个层面的优化。本篇就是我个人现场调试与部署的实战总结，涵盖了完整的技术细节、硬件配置、评测数据和代码示例。

---

一、系统与硬件环境概览

在开始之前，下面是我们用于构建分布式训练平台的香港GPU服务器www.a5idc.com规格表（3节点）：

组件	节点A	节点B	节点C
操作系统	CentOS 8.6 (Kernel 4.18)	CentOS 8.6	CentOS 8.6
CPU	2× Intel Xeon Gold 6248 (40 cores)	同左	同左
内存	512 GB DDR4	512 GB DDR4	512 GB DDR4
GPU	4× NVIDIA A100 80GB	4× NVIDIA A100 80GB	4× NVIDIA A100 80GB
GPU互联	NVLink 3.0	NVLink 3.0	NVLink 3.0
网络	Mellanox 100GbE RoCE V2	Mellanox 100GbE RoCE V2	Mellanox 100GbE RoCE V2
存储	4× 2TB NVMe	4× 2TB NVMe	4× 2TB NVMe

关键指标说明

• GPU型号：A100 80GB是面向大规模训练的高端加速卡，具备Tensor Core和更大显存。
• 互联与网络：NVLink用于节点内GPU互联，高带宽低延迟；Mellanox 100GbE RoCE V2用于节点间通信。
• 操作系统：CentOS 8官方标准库对一些新软件支持有限，因此很多组件需要从源或第三方构建。

---

二、基础依赖安装与系统调优

2.1 安装NVIDIA驱动与CUDA

我们选择CUDA 12.1，与PyTorch 2.1兼容性良好。步骤如下：

2.1.1 安装驱动

# 禁用nouveau驱动
cat <<EOF > /etc/modprobe.d/blacklist-nouveau.conf
blacklist nouveau
options nouveau modeset=0
EOF
dracut --force

# 重启
reboot

# 添加NVIDIA仓库
dnf config-manager --add-repo=https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo

# 安装驱动和CUDA
dnf clean all
dnf -y module disable nvidia-driver
dnf -y install nvidia-driver-latest-dkms cuda-12-1

安装完成后，确认驱动与GPU状态：

nvidia-smi

2.2 安装cuDNN

从NVIDIA官网获取cuDNN 8.9 for CUDA 12.1 RPM包，然后：

dnf -y install libcudnn8 libcudnn8-devel

2.3 安装Mellanox OFED与RDMA

为了实现RoCE V2，我们必须安装并启用Mellanox OFED：

# 下载与节点网络硬件匹配的MLNX_OFED
bash MLNX_OFED_LINUX-<version>-<build>.run

# 启用RDMA
systemctl enable --now rdma

验证RDMA接口：

rdma link show

2.4 系统网络优化

编辑 /etc/sysctl.conf:

net.core.rmem_max=268435456
net.core.wmem_max=268435456
net.ipv4.tcp_rmem=4096 87380 268435456
net.ipv4.tcp_wmem=4096 65536 268435456
net.ipv4.tcp_congestion_control=bbr

使设置生效：

sysctl -p

---

三、PyTorch 与分布式通信库安装

3.1 安装Miniconda

为了隔离环境，我们通过Miniconda部署：

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

创建环境：

conda create -n torch_dist python=3.10
conda activate torch_dist

3.2 安装PyTorch及NCCL

PyTorch 2.1具有完善的分布式支持：

# 指定CUDA 12.1
pip install torch==2.1.0+cu121 torchvision==0.17.0+cu121 \
    --extra-index-url https://download.pytorch.org/whl/cu121

安装NCCL 2.18（可选，但推荐用于高效通信）：

yum -y install nccl-2.18.3-1+cuda12.1

确认NCCL版本：

rpm -qa | grep nccl

---

四、实现分布式训练（以ResNet50为例）

4.1 分布式训练思想

PyTorch分布式训练核心是torch.distributed模块，它允许在多个进程间共享梯度与参数。主流后端包括：

后端	说明	推荐场景
NCCL	适用于GPU	多GPU训练首选
GLOO	纯CPU，网络容错较好	CPU训练
MPI	基于MPI库	混合场景与老旧集群

我们使用NCCL做为后端，并辅以TorchElastic进行弹性节点管理。

4.2 配置SSH信任

在三台节点之间设置SSH免密：

ssh-keygen -t rsa -b 4096
# 分发到各节点
ssh-copy-id user@nodeB
ssh-copy-id user@nodeC

4.3 启动分布式训练脚本

创建名为 train_dist.py：

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models

def main():
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)

    # Data
    transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.ToTensor(),
    ])
    dataset = datasets.FakeData(transform=transform)
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=64, sampler=train_sampler, num_workers=8)

    # Model
    model = models.resnet50().cuda()
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    # Optimizer
    optimizer = optim.SGD(model.parameters(), lr=0.1)

    # Training
    model.train()
    for epoch in range(5):
        train_sampler.set_epoch(epoch)
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.cuda(), target.cuda()
            optimizer.zero_grad()
            output = model(data)
            loss = nn.CrossEntropyLoss()(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 10 == 0 and local_rank == 0:
                print(f'Epoch {epoch} Batch {batch_idx} Loss {loss.item()}')

if __name__ == "__main__":
    main()

4.4 启动命令

使用Torch Elastic机制：

python -m torch.distributed.launch \
    --nproc_per_node=4 \
    --nnodes=3 \
    --node_rank=0 \
    --master_addr="10.0.0.1" \
    --master_port=29500 \
    train_dist.py

在节点B、C上分别设置node_rank=1和node_rank=2即可。

---

五、性能调优与实测结果

5.1 NCCL参数优化

为了最大化网络带宽：

export NCCL_DEBUG=INFO
export NCCL_IB_HCA=mlx5_0
export NCCL_IB_GID_INDEX=3
export NCCL_IB_SL=5
export NCCL_TCP_KEEPALIVE=1

5.2 参数对比测试

我们对比了以下训练设置下的每秒样本处理量（images/sec）：

配置	单卡	4卡同节点	12卡三节点
默认（无优化）	112	376	960
+RDMA RoCE V2	112	428	1150
+NCCL调优	112	450	1250

分析：

• 单节点4卡性能接近线性扩展，但节点间通信成为瓶颈；
• 启用RoCE V2后网络带宽利用率提高；
• NCCL参数调优结合RDMA环境，整体12卡性能提升近30%。

---

六、常见问题与解决方案

6.1 节点间延迟高

检查RoCE配置：

ibv_devinfo

确保QoS与MTU一致（建议MTU 9000）。

6.2 NCCL超时失败

缩短超时时间：

export NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_TIMEOUT=300

6.3 数据加载成为瓶颈

提升DataLoader：

DataLoader(..., num_workers=16, pin_memory=True)

---

七、总结与推荐实践

A5数据通过本次实战，在CentOS 8服务器集群上完成了PyTorch分布式训练的部署与优化，关键实践结论如下：

• GPU驱动、CUDA和cuDNN版本必须严格匹配；
• RDMA RoCE V2 + NCCL是节点间通信的最佳组合；
• PyTorch DistributedDataParallel比单进程多卡更稳定；
• 网络与内存参数优化对性能提升明显。

最终，我们在推荐模型训练任务中，将单卡训练时间从16小时缩短至约1小时（12卡分布式）。这一提升直接缩短了模型迭代周期，为业务提供了明显效率优势。