集群中多节点NCCL分布式通信

集群中多节点NCCL分布式通信

需求：在多台linux服务器中，我需要配置分布式通信环境，以适配后续分布式训练大模型的需求。

测试脚本：

import torch
import torch.distributed as dist
from datetime import timedelta

def demo_basic():
    dist.init_process_group("nccl", init_method='env://', timeout=timedelta(seconds=20))
    rank = dist.get_rank()
    print(f"Start running on rank {rank}.")
    
    torch.cuda.set_device(rank % torch.cuda.device_count())

    device_id = rank % torch.cuda.device_count()

    tensor = torch.arange(2) + 1 + 2 * rank
    tensor = tensor.cuda()
    print(f"rank {rank}: {tensor}")
    
    tensor_list = [torch.zeros_like(tensor).cuda() for _ in range(4)]
    dist.all_gather(tensor_list, tensor)
    print(f"after gather, rank {rank}: tensor_list: {tensor_list}")
    
    dist.barrier()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f"after reduce, rank {rank}: tensor: {tensor}")

if __name__ == "__main__":
    demo_basic()

我有两台服务器，每台服务器中有2张GPU，两台服务器能相互免密登录，在终端中运行脚本：

torchrun --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=172.18.xxx.xx --master_port=29500 mptest.py

在不同服务器上需更改–node_rank=0的值。

服务器0输出：

服务器1输出：

踩坑

若服务器有多个网卡，需要在代码运行前在终端指定网卡，不然会报错！

export NCCL_SOCKET_IFNAME=eno2

每台服务器都需要！且每台服务器的网卡名不同，可通过ipconfig进行查看！

若仍出现无法连接的情况，可以参考：
https://blog.youkuaiyun.com/Komach/article/details/130765773