【CS336】lecture7 多卡|通信架构|集合通信|NCCL|通信原语|torch.distributed

多GPU通信架构

这讲开始，进入多GPU并行。整体结构如上，我们前几讲了解的是单GPU的结构，在此之上还有更高的层次：单节点多GPU之间通过NVLink和NVSwitch连接，这也被称为单机多卡。多节点之间通过以太网或Infiniband连接，这被称为多机多卡

根据我们之前单GPU的内存层次学到的规律，离得越近通信应该越快，在这里同样满足这个规律，NVLink比HBM慢，以太网又比NVLink慢

使用多GPU的是模型参数太大，单卡显存放不下，这是从内存角度来讲的。另一方面，多卡可以获得更高的计算吞吐，这是从计算角度来讲的。

基础概念

分布式里最基础的操作原语集合集合被称为Collective operation，集合通信，这是很早就定义了的分布式通信原语。原语指的就是一些概念性的基础操作，不规定具体实现，而是定义这个操作要实现什么，举个例子，并发编程里的P/V操作，或者说wait/signal就是一种原语。

基于集合通信的通信库有很多，最经典的就是NCCL（NVIDIA Collective Communications Library），英伟达集合通信库，负责英伟达GPU上的集合通信。

rank指的就是设备编号，在这里主要指GPU编号。world size指设备总数，在这图里就是4

基本操作

介绍一下集合通信的基本操作，后面的并行策略都要基于这些基础操作

Broadcast

从rank0原样复制到其他rank

py 复制代码

# Input
rank0 = tensor([0., 1, 2, 3])

# Output
rank0 = tensor([0., 1, 2, 3])
rank1 = tensor([0., 1, 2, 3])
rank2 = tensor([0., 1, 2, 3])
rank3 = tensor([0., 1, 2, 3])

Scatter

意思是分散，注意和广播的区别，是把一个列表均分成多份发送，不是直接发送。这里一个一维张量有四个元素，均分四份就是每个rank分到一个元素，准确来说是只含一个元素的列表，仍然是个列表。

py 复制代码

# Input
rank0 = tensor([0., 1, 2, 3])

# Output
rank0 = tensor([0.])
rank1 = tensor([1.])
rank2 = tensor([2.])
rank3 = tensor([3.])

Gather

scatter的逆运算，不同rank的数据都发到一个rank组合起来

py 复制代码

# Input
rank0 = tensor([0.])
rank1 = tensor([1.])
rank2 = tensor([2.])
rank3 = tensor([3.])

# Output
rank0 = tensor([0., 1, 2, 3])

Reduce

这个更常见一点，对所有rank的数据应用某个操作，比如max,sum,min，在rank0得到一个结果

py 复制代码

# Input
rank0 = tensor([0.])
rank1 = tensor([1.])
rank2 = tensor([2.])
rank3 = tensor([3.])

# Output
rank0 = tensor([6.])  # Sum of all ranks (0 + 1 + 2 + 3)

All-gather

收集到所有rank，朴素的实现思路是先gather到一个rank，再broadcast到所有rank，但是broadcast是个低效的操作。高级实现是ring gather，环形收集，这里不展开

py 复制代码

rank0 = tensor([0.])
rank1 = tensor([1.])
rank2 = tensor([2.])
rank3 = tensor([3.])

# Output
rank0 = tensor([0., 1, 2, 3])
rank1 = tensor([0., 1, 2, 3])
rank2 = tensor([0., 1, 2, 3])
rank3 = tensor([0., 1, 2, 3])

Reduce-scatter

在每个维度上分别reduce，然后把不同维度scatter分给不同rank。朴素的实现是先reduce到一个rank上，再scatter到不同rank。和前面的all-gather类似，从一个rank scatter到所有rank是比较低效的。生产级实现也不会这么做。

py 复制代码

# Input
rank0 = tensor([0., 1, 2, 3])
rank1 = tensor([1., 2, 3, 4])
rank2 = tensor([2., 3, 4, 5])
rank3 = tensor([3., 4, 5, 6])

# Output
rank0 = tensor([6.])   # Sum along dim 0 (0 + 1 + 2 + 3)
rank1 = tensor([10.])  # Sum along dim 1 (1 + 2 + 3 + 4)
rank2 = tensor([14.])  # Sum along dim 2 (2 + 3 + 4 + 5)
rank3 = tensor([18.])  # Sum along dim 3 (3 + 4 + 5 + 6)

All-reduce

All-reduce = reduce-scatter + all-gather

reduce之后的结果要发给所有rank，最低效的实现和reduce scatter类似，先reduce到一个rank，再广播给所有rank。利用我们前面定义过的操作，高效一些的思路是先scatter-reduce，让每个rank分别持有一个维度的reduce结果，再all-gather收集到每个rank。这样的瓶颈在于scatter-reduce和all-gather的实现，把这俩实现好了就可以很高效。

py 复制代码

# Input
rank0 = tensor([0., 1, 2, 3])
rank1 = tensor([1., 2, 3, 4])
rank2 = tensor([2., 3, 4, 5])
rank3 = tensor([3., 4, 5, 6])

# Output
rank0 = tensor([6., 10, 14, 18])
rank1 = tensor([6., 10, 14, 18])
rank2 = tensor([6., 10, 14, 18])
rank3 = tensor([6., 10, 14, 18])

All-to-all

开销最大的操作，任意两个rank之间都有通信

下面的例子只是演示，不一定要传这些。下面的具体思路是每个rank都把自己的第i维信息传给rank i

py 复制代码

# Input
rank0 = tensor([0., 1, 2, 3])      # send  0 to rank 0,  1 to rank 1,  2 to rank 2,  3 to rank 3
rank1 = tensor([4., 5, 6, 7])      # send  4 to rank 0,  5 to rank 1,  6 to rank 2,  7 to rank 3
rank2 = tensor([8., 9, 10, 11])    # send  8 to rank 0,  9 to rank 1, 10 to rank 2, 11 to rank 3
rank3 = tensor([12., 13, 14, 15])  # send 12 to rank 0, 13 to rank 1, 14 to rank 2, 15 to rank 3

# Output
rank0 = tensor([0, 4, 8, 12])
rank1 = tensor([1, 5, 9, 13])
rank2 = tensor([2, 6, 10, 14])
rank3 = tensor([3, 7, 11, 15])

torch.distributed

torch库也有分布式库，和torch定位类似，不追求最高级的性能，但是提供的要保证正确性和基本的性能，用于快速搭建，验证流程。

下面给一个torch分布式库验证all-reduce=scatter-reduce+all-gather的代码，使用的数据就是上面all-reduce的数据，4个rank，每个保存一个长度4的一维张量

分别调用torch的all-reduce接口，和scatter-reduce+all-gather接口，分别打印结果进行对比。结果应该是相同的。

py 复制代码

def collective_operations_main(rank: int, world_size: int):
    setup(rank, world_size)

    # All-reduce
    dist.barrier()
    data = tensor(
        [0., 1, 2, 3], device=cuda_if_available(rank)
    ) + rank

    print(f"Rank {rank} [before all-reduce]: {data}", flush=True)
    dist.all_reduce(
        tensor=data, op=dist.ReduceOp.SUM, async_op=False
    )
    print(f"Rank {rank} [after all-reduce]: {data}", flush=True)

    # Reduce-scatter
    dist.barrier()
    input = torch.arange(
        world_size, dtype=torch.float32,
        device=cuda_if_available(rank),
    ) + rank
    output = torch.empty(1, device=cuda_if_available(rank))

    print(
        f"Rank {rank} [before reduce-scatter]: "
        f"input = {input}, output = {output}", flush=True,
    )
    dist.reduce_scatter_tensor(
        output=output, input=input,
        op=dist.ReduceOp.SUM, async_op=False,
    )
    print(
        f"Rank {rank} [after reduce-scatter]: "
        f"input = {input}, output = {output}", flush=True,
    )

    # All-gather
    dist.barrier()
    input = output
    output = torch.empty(world_size, device=cuda_if_available(rank))

    print(
        f"Rank {rank} [before all-gather]: "
        f"input = {input}, output = {output}", flush=True,
    )
    dist.all_gather_into_tensor(
        output_tensor=output, input_tensor=input, async_op=False,
    )
    print(
        f"Rank {rank} [after all-gather]: "
        f"input = {input}, output = {output}", flush=True,
    )

    # Indeed, all-reduce = reduce-scatter + all-gather!
    cleanup()