when using multiple GPUs, `loss.mean()` may have subtle bias

The problem is that, when the input is distributed to multiple GPUs, the input on each GPU may have different batch_size.

For example, if you have 2 GPUs and the total batch_size is 13, then the batch_size for each GPU will be 7 and 6 respectively, loss.mean() will not give the exact loss. Although it may have little influence on the training of the model, it is not the exact result.

https://github.com/huggingface/pytorch-pretrained-BERT/blob/3fc63f126ddf883ba9659f13ec046c3639db7b7e/examples/run_squad.py#L1006-L1007

[yaroslavvb]

You could fix it with something like this

batch_size = torch.tensor(data.shape[1]).to(device)
dist.all_reduce(batch_size, op=dist.ReduceOp.SUM)
dist.all_reduce(loss, op=dist.ReduceOp.SUM)
mean_loss = loss/batch_size

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.