Network-Offloaded Bandwidth-Optimal Broadcast and Allgather for Distributed AI
<|reference_start|>Network-Offloaded Bandwidth-Optimal Broadcast and Allgather for Distributed AI: In the Fully Sharded Data Parallel (FSDP) training pipeline, collective operations can be interleaved to maximize the communication/computation overlap. In this scenario, outstanding operations such as Allgather and Reduce-Scatter can compete for the injection bandwidth and create pipeline bubbles. To address this problem, we propose a novel bandwidth-optimal Allgather collective algorithm that leverages hardware multicast. We use multicast to build a constant-time reliable Broadcast protocol, a building block for constructing an optimal Allgather schedule. Our Allgather algorithm achieves 2x traffic reduction on a 188-node testbed. To free the host side from running the protocol, we employ SmartNIC offloading. We extract the parallelism in our Allgather algorithm and map it to a SmartNIC specialized for hiding the cost of data movement. We show that our SmartNIC-offloaded collective progress engine can scale to the next generation of 1.6 Tbit/s links.<|reference_end|>
@article{khalilov2024network-offloaded,
    title={Network-Offloaded Bandwidth-Optimal Broadcast and Allgather for
  Distributed AI},
    author={Mikhail Khalilov, Salvatore Di Girolamo, Marcin Chrapek, Rami
  Nudelman, Gil Bloch, Torsten Hoefler},
    journal={arXiv preprint arXiv:2408.13356},
    year={2024},
    archivePrefix={arXiv},
    eprint={2408.13356},
    primaryClass={cs.DC}
}