We use an all_reduce operation on the gradient, which sums them, because DDP
will divide them later by the number of devices. One can use the operation by simply
calling: FullGatherLayer.apply(x) on the input x. Practically, for the methods above,
this needs to be done on the embeddings just before the computation of the loss.

Additional tricks. We advise to always use the eﬀective batch size as argument to the
training script, as well as for comparing runs. The DataLoader class takes the per device
batch size as, argument, which can be obtained by dividing the eﬀective batch size by the
number of devices which is world_size in PyTorch. We also advise to use an adaptive learn-

35

"""
Gather tensors from all process and support backward propagation
for the gradients across processes.
"""

Algorithm 1:
1 class GatherLayer(torch.autograd.Function):
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
A Cookbook of Self-Supervised Learning