ReduceScatter
Reduce-Scatter: perform reduce on each element, scatter results.
# 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)https://tech.preferred.jp/en/blog/technologies-behind-distributed-deep-learning-allreduce/ is really good
Use case?
After backward pass, sum gradients from different data shards, but distribute storage for FSDP.
- Reduce-scatter combines the gradients while partitioning the result
- all-gather then reconstructs the full reduced gradient on every rank.
- In practice it’s just an All-Reduce, but being able to split it up this way allows for more efficient implementation