Collective Operations

Broadcast (Collective Operation)

Learned from CS336. There’s also PyTorch Broadcasting which is a similar idea, but used in a different context.

Broadcast: copy from rank 0 to all ranks

# 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])

What is the use-case?

For example, when you are initially loading the model weights on rank0, you will broadcast the weights to all the other ranks with DDP.

• https://docs.pytorch.org/tutorials/intermediate/ddp_tutorial.html “DDP broadcasts model states from rank 0 process to all other processes in the DDP constructor”
• Because init_sync=True by default, that is the flag that does that