How to scale learning rate with batch size for DDP training?
When using LARS optimizer, usually the batch size is scale linearly with the learning rate.
Suppose I set the base_lr to be 0.1 * batch_size / 256.
Now for 1 GPU training with batch size 512, the learning rate should be 0.1 * 2 = 0.2
However when I use 2 GPUs with DDP backend and batch size of 512 on each GPU. Should my learning rate be:

0.1 * 2 = 0.2
or 0.1 * 2 *  2 (no. GPUs) = 0.4
As far as I know, learning rate is scaled with the batch size so that the sample variance of the gradients is kept approx. constant.
Since DDP averages the gradients from all the devices, I think the LR should be scaled in proportion to the effective batch size, namely, batch_size * num_accumulated_batches * num_gpus * num_nodes
In this case, assuming batch_size=512, num_accumulated_batches=1, num_gpus=2 and num_noeds=1  the effective batch size is 1024, thus the LR should be scaled by sqrt(2), compared to a single gpus with effective batch size 512.