Correctness of "counter average weight reduction"

[youyuandx]

Hi devs, first, thanks for this clean & useful code base!

I'm not following this line

tevatron/src/tevatron/retriever/modeling/encoder.py

Line 72 in f88e0c7

loss = loss * self.world_size # counter average weight reduction

Since in an earlier line

tevatron/src/tevatron/retriever/modeling/encoder.py

Line 41 in f88e0c7

self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')

reduction is "mean", I believe it is correct to let DDP average the gradient across multiple GPUs, and there's no need to counteract that? Only when DDP does the average, would it be equivalent between 2 GPUs with 8 examples each and 1 GPU with 16 examples, both reduced by "mean".

Thanks in advance.

[MXueguang]

Hi @rgdyy , #30 is relevant.

[youyuandx]

Hi @MXueguang , sorry that I forgot about this for a long while. As far as I understood, mean of mean is still correctly, mean, so Line 72 loss * world_size should be deleted. What do you think?

[youyuandx]

It's been a long time again. I read #30 and also another thread. After skimming the repo pointed to from there, I believe the loss scaling is not sufficient: because the gather is non-diff, cross-GPU negatives would receive no grad from the current query's contrastive loss. However, I do not have the time to read and validate all things in that repo, so I would personally choose the solution discussed here with PyTorch Lightning. Please let me know if I'm wrong. Thanks for the code and everything!

[MXueguang]

Thanks for bring this back. I'll carefully take a look the implementations and the affects in model performance recently and let you know @rgdyy

[youyuandx]

Thanks @MXueguang ! Also hope to share the not-quite-official APIs in PyTorch with grad propagation: https://discuss.pytorch.org/t/do-gradients-propagate-through-all-reduce-all-gather/159572
Also I understood from this issue that pytorch-lightning actually uses these PyTorch APIs. The issue is still open but I understood it's already switched to these APIs but the issue was not cleaned up.