PyTorch Distributed 与 Collective Communication

torch distributed

torch distributed以进程组(process group, PG) 作为整个通信系统，每个PG有自己的world_size与成员process，每个process拥有自己的rank，以此为基础进行通信。通信后端支持nccl(n卡GPU), gloo(CPU), MPI(各种Message Passing Interface是西安)以及自定义PG后端, torch 2.x还支持ucc(Nvidia推出的统一collective通信库, xccl(XPU Collective Communication Library, 围绕Intel GPU/XPU).

torch distributed通常需要定义一些进程内环境变量：

• MASTER_ADDR: 表示RANK0所在机器的地址(IP/hostname)
• WORLD_SIZE: PG成员数量
• RANK: 当前process的全局ID(0~WORLD_SIZE-1)
• LOCAL_RANK: process在本机内的编号

LOCAL_RANK 通常用来设置初始时torch cuda设置device，比如:

local_rank=init(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

torch.distributed里和PG相关的操作主要有:

• init_process_group : 初始化默认PG
• destroy_process_group: 摧毁释放指定PG,未指定则释放默认PG
• new_group : 在默认PG的基础上，按照指定ranks创建子PG
• is_initialized: 默认PG是否被初始化
• get_backend: 获取指定PG的backend
• get_world_size:获取指定PG的world_size
• MASTER_PORT: RANK0上的一个空闲窗口

与communication相关的操作主要有:

• torch.distributed.broadcast(tensor, src, group=None)
• torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=None)
• torch.distributed.reduce(tensor, dst, op=..., group=None)
• torch.distributed.all_gather(tensor_list, tensor, group=None)
• torch.distributed.gather(tensor, dst, group=None)
• torch.distributed.scatter(tensor, scatter_list, src, group=None)
• torch.distributed.reduce_scatter(output, input_list, op=..., group=None)
• torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None)
• torch.distributed.all_gather_object(object_list, obj, group=None)
• torch.distributed.broadcast_object_list(obj_list, src, group=None)

除去上述这8中collective(还有一些和python object相关)，还支持P2P通信:

• torch.distributed.send(tensor, dst)
• torch.distributed.recv(tensor, src)

torchrun

torchrun 用来启动多进程python脚本，适用与单机多卡/多机多卡，比如:

# Machine 0
torchrun --nprocs_per_node=4 --nnodes=2 --node_rank=0 --master_addr=x  \\
	--master_port=x test.py
# Machine 1
torchrun --nprocs_per_node=4 --nnodes=2 --node_rank=1 --master_addr=x  \\
	--master_port=x test.py

torchrun 会自动为每个process创建一些环境变量,包括之前提到的torch distributed中使用的五个环境变量，还有 PYTHON_EXEC (所使用的python解释器）, LOCAL_WORLD_SIZE (本机进程数）

collective communication

主要包括: Gather, Reduce, Scatter, AllReduce, AllGather, ReduceScatter, Broadcast, AllToAll。

NCCL 文档: Collective Operations — NCCL 2.28.9 documentation

Broadcast

将某个并行组成员持有的数据传输给其他并行组成员。

import torch
import torch.distributed as dist
dist.init_process_group(backend="nccl", rank=0,world_size=4)
tensor = torch.randn(3,3)
dist.broadcast(tensor,src=0)

Gather

PG中每process持有同大小的数据，将所有数据汇集在某一RANK process上。

    tensor = torch.zeros(world_size).cuda()
    for i in range(world_size):
        tensor[i] = rank * world_size + i

    if rank == 0:
        tensor_list = [torch.zeros(world_size).cuda() for _ in range(world_size)]
    else:
        tensor_list = None
    dist.gather(tensor, gather_list=tensor_list, dst=0)

每RANK传入自己的数据，并且DST RANK需要准备接收Buffer(示例中的 gather_list)

Reduce

所有并行组成员都持有一份同大小的数据，执行Reduce规约操作(SUM求和，PROD逐元素相乘，MAX取最大值，MIN取最小值,AVG求平均等等，以及自定义CustomOp),其结果保存在某一rank成员中。

    reduce_op = dist.ReduceOp.SUM
    tensor = torch.tensor([local_rank],device=f"cuda:{local_rank}")
    res = dist.reduce(tensor=tensor, dst=1, op=reduce_op)

Scatter

PG中的某一RANK将自身持有的N份等大小数据按RANK分给每一个RANK。

    if local_rank == 0:
        tensor_list = [
            torch.tensor([float(i)], dtype=torch.float32, device="cuda")
            for i in range(world_size)
        ]
    else:
        tensor_list = None
    tensor = torch.zeros(1, dtype=torch.float32).cuda()
    dist.scatter(tensor, scatter_list=tensor_list, src=0)

dist.scatter 会将scatter_list(list of tensor，大小为PG的world_size)中的数据按rank分散给每个rank.

AllReduce

PG中的每个等大小数据规约后将结果分散给每个rank,等价于reduce+broadcast。

		reduce_op = dist.ReduceOp.SUM
    tensor = torch.tensor([local_rank],device=f"cuda:{local_rank}")
    res = dist.all_reduce(tensor, op=reduce_op)

AllGather

PG中每个rank持有同大小数据，将gather后的数据传输给每个rank,等价于gather+broadcast

    world_size = dist.get_world_size()
    tensor_list = [torch.zeros(1).cuda() for _ in range(world_size)]
    dist.all_gather(tensor_list, torch.tensor([float(local_rank)]).cuda())

ReduceScatter

对于一组input_list of tensor,进行reduce之后scatter到每个rank

    input_list = []
    for i in range(world_size):
        t = torch.full((2,), rank * 10 + i, device="cuda", dtype=torch.float32)
        input_list.append(t)

    output = torch.empty(2, device="cuda", dtype=torch.float32)

    dist.reduce_scatter(
        output=output,
        input_list=input_list,
        op=dist.ReduceOp.SUM
    )

input_list中元素得与output 大小相同

AllToAll

PG里每个rank都保留world_size * world_size 块等大小的buffer, 分别将buffer发送到对应的rank，从rank接收对应的buffer.

    input_list = []
    for i in range(world_size):
        t = torch.full((2,), rank * 10 + i,
                       device="cuda", dtype=torch.float32)
        input_list.append(t)

    # 准备 output_list（和 input_list 同 shape）
    output_list = [torch.empty_like(input_list[0])
                   for _ in range(world_size)]

    # all_to_all：每个 rank 把 input_list[i] 发给 rank i
    dist.all_to_all(output_list, input_list)