RDMA：使用高性能网络进行分布式训练

在大模型的AI并行计算中，通过降低通信量、计算和通信交叠、提升通信效率来优化计算性能。本文介绍如何在Scnet使用高性能网络进行配置，以实现上述目标。

一、配置高性能网络变量

Scnet模型训练已开启RDMA特性并设置了最优NCCL变量。建议使用平台预置的默认变量以获得较优性能。您也可以根据训练框架、通信框架以及模型特点进行必要的调整。

默认变量

平台已在系统中预置了默认变量，具体请参照下表内容：

export NCCL_IB_TC=136
export NCCL_IB_GID_INDEX=3
export NCCL_SOCKET_IFNAME=eth
export NCCL_DEBUG=INFO
export NCCL_IB_HCA=mlx5
export NCCL_IB_TIMEOUT=22
export NCCL_IB_QPS_PER_CONNECTION=8
export NCCL_IB_DISABLE=0

默认环境变量说明：

关键环境变量	环境变量说明
NCCL_IB_TC	匹配网络映射规则，未配置或配置错误可能导致性能受损。
NCCL_IB_GID_INDEX	提供推荐值，未配置或配置错误会导致NCCL报错。
NCCL_SOCKET_IFNAME	选择正确的端口以建立连接，不同规格对端口的要求不同。未配置或配置错误可能会导致NCCL建连失败。
NCCL_DEBUG	通常将日志级别设置为INFO，以获得更多NCCL相关的日志输出，帮助定位和解决潜在问题。
NCCL_IB_HCA	指定RDMA通信的网卡。不同算力节点下IBdev的数量和命名规则有所区别，未配置或配置错误可能导致性能受损。
NCCL_IB_TIMEOUT	增加RDMA连接超时时间，提升训练任务的容错性能，未配置或配置错误可能导致训练任务中断。
NCCL_IB_QPS_PER_CONNECTION	适当增加每个连接的QP（Queue Pair）数量，有效提高网络吞吐率。
NCCL_IB_DISABLE	用于控制是否禁用 InfiniBand (IB) 通信。

二、RDMA使用案例

PyTorch

步骤 1: 环境准备

1. 硬件/网络要求：
   

• 确保两台节点通过支持 RDMA 的网络（如 InfiniBand 或 RoCE）连接。
  
• 确认 ibstatus 或 ibv_devices 命令能检测到 RDMA 设备。
  

2. 软件依赖：
   

• 安装相同版本的 PyTorch 和 CUDA。
  
  • 安装 NCCL 并启用 RDMA 支持（通常包含在 PyTorch 的 CUDA 版本中）。
    
• 安装 RDMA 驱动（如 libibverbs-dev）。
  

步骤 2: 编写 PyTorch 分布式代码

# File: nccl_rdma_example.py
import os
import torch
import torch.distributed as dist
from torch.multiprocessing import Process

def run(rank, world_size):
    # 设置当前进程的 CUDA 设备（假设每个节点只有1个GPU）
    torch.cuda.set_device(rank)
    
    # 初始化进程组
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        rank=rank,
        world_size=world_size
    )
    
    # 创建一个张量并传输到GPU
    tensor = torch.ones(1024, device=f'cuda:{rank}')
    
    # 执行All-Reduce操作
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    
    print(f"Rank {rank} received result: {tensor[:5]} ...")

if __name__ == "__main__":
    # 从环境变量获取参数
    world_size = int(os.environ['WORLD_SIZE'])
    rank = int(os.environ['RANK'])
    
    run(rank, world_size)

步骤 3: 启动脚本

在节点1 (172.20.170.60) 上执行：

export MASTER_ADDR=172.20.170.63
export MASTER_PORT=29500  # 选择一个未被占用的端口
export WORLD_SIZE=2
export RANK=0

# 可选：强制启用NCCL的IB/RDMA传输，如果在vcjob创建pod内，下面环境变量已经自动注入
export NCCL_IB_DISABLE=0  # 启用InfiniBand
export NCCL_SOCKET_IFNAME=ib0  # 指定RDMA网络接口名,如果在pod内是eth0
export NCCL_DEBUG=INFO  # 查看调试信息

python nccl_rdma_example.py

在节点2 (172.20.0.195) 上执行：

export MASTER_ADDR=172.20.170.63
export MASTER_PORT=29500
export WORLD_SIZE=2
export RANK=1

# 同样设置NCCL环境变量，如果在vcjob创建pod内，下面环境变量已经自动注入
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=ib0 # 如果在pod内是eth0
export NCCL_DEBUG=INFO

python nccl_rdma_example.py

验证 RDMA 使用 检查 NCCL 的调试输出：

NCCL INFO Connected all rings
NCCL INFO Connected tree
NCCL INFO Channel 00 : 0[9000] 1[9000] via P2P/IPC
NCCL INFO Channel 01 : 0[9000] 1[9000] via NET/IB/0

出现 via NET/IB/0 表示成功使用 RDMA。