多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
# Inputrank0=tensor([0.,1,2,3])# Outputrank0=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分到一个元素，准确来说是只含一个元素的列表，仍然是个列表。
# Inputrank0=tensor([0.,1,2,3])# Outputrank0=tensor([0.])rank1=tensor([1.])rank2=tensor([2.])rank3=tensor([3.])Gather
scatter的逆运算，不同rank的数据都发到一个rank组合起来
# Inputrank0=tensor([0.])rank1=tensor([1.])rank2=tensor([2.])rank3=tensor([3.])# Outputrank0=tensor([0.,1,2,3])Reduce
这个更常见一点，对所有rank的数据应用某个操作，比如max,sum,min，在rank0得到一个结果
# Inputrank0=tensor([0.])rank1=tensor([1.])rank2=tensor([2.])rank3=tensor([3.])# Outputrank0=tensor([6.])# Sum of all ranks (0 + 1 + 2 + 3)All-gather
收集到所有rank，朴素的实现思路是先gather到一个rank，再broadcast到所有rank，但是broadcast是个低效的操作。高级实现是ring gather，环形收集，这里不展开
rank0=tensor([0.])rank1=tensor([1.])rank2=tensor([2.])rank3=tensor([3.])# Outputrank0=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是比较低效的。生产级实现也不会这么做。
# Inputrank0=tensor([0.,1,2,3])rank1=tensor([1.,2,3,4])rank2=tensor([2.,3,4,5])rank3=tensor([3.,4,5,6])# Outputrank0=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的实现，把这俩实现好了就可以很高效。
# Inputrank0=tensor([0.,1,2,3])rank1=tensor([1.,2,3,4])rank2=tensor([2.,3,4,5])rank3=tensor([3.,4,5,6])# Outputrank0=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
# Inputrank0=tensor([0.,1,2,3])# send 0 to rank 0, 1 to rank 1, 2 to rank 2, 3 to rank 3rank1=tensor([4.,5,6,7])# send 4 to rank 0, 5 to rank 1, 6 to rank 2, 7 to rank 3rank2=tensor([8.,9,10,11])# send 8 to rank 0, 9 to rank 1, 10 to rank 2, 11 to rank 3rank3=tensor([12.,13,14,15])# send 12 to rank 0, 13 to rank 1, 14 to rank 2, 15 to rank 3# Outputrank0=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接口，分别打印结果进行对比。结果应该是相同的。
defcollective_operations_main(rank:int,world_size:int):setup(rank,world_size)# All-reducedist.barrier()data=tensor([0.,1,2,3],device=cuda_if_available(rank))+rankprint(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-scatterdist.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-gatherdist.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()
