When using torch.nn.parallel.data_parallel for distributed training, models are copied onto multiple GPU’s and can complete a forward pass without the copied models effecting each other. However, how do the copied models interact during the backward pass? How are the model weights updated on each GPU?
When reading the documentation [1] I see the explanation:
“gradients from each replica are summed into the original module”
but I’m not sure how to interpret this statement. My best guess is that each GPU’s model is updated with the sum of all GPU’s model’s gradients, which I would then interpret that there is locking across GPU’s so they each start training on a new mini-batch only after they all finish processing their current mini-batch.
[1] https://pytorch.org/docs/master/nn.html?highlight=dataparallel#dataparallel-layers-multi-gpu-distributed 41