AllGather
All-gather: perform gather to all ranks, not just rank 0
# Input
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])Use case?
For FSDP, each rank holds parameter shard, gather to get full parameters for forward pass
Implementation
Uses the ring algorithm.
Another way is to use bruck’s algorithm