Suppose I have updated the neural network weights via some operation on the master node (rank = 0), and I would like to broadcast updated weights to all other workers in the group.
In the torch.distributed API I see several functions:
- torch.distributed.broadcast
- torch.distributed.broadcast_object_list
- torch.distributed.broadcast_multigpu
Seems, like a simple option:
for param in model.parameters():
dist.broadcast(param, src=0)
Would work. The third option seems to be more specific for training on single node Multi GPU but it is not clear how?
What is the recommended way to broadcast parameters for NN?
If it is done prior to the training loop, does the default option async_op=False mean, that other processes would not start execution their part of the loop until the main process broadcasts the weights?