`all_reduce` does not apply `scale` when `xr.world_size == 1`

[afzalxo]

❓ Questions and Help

Hi, I have noticed that when world_size == 1, all_reduce is a no-op and does not apply scale:

In torch_xla.core.xla_model in def all_reduce:

# No-op if there is only one device
  if runtime.world_size() == 1 and not xu.getenv_as('XLA_ALWAYS_ALLREDUCE',
                                                    bool, False):
    if isinstance(inputs, torch.Tensor):
      return inputs.clone()
    else:
      return inputs

Is this intended behavior? If it is indeed intended, it makes the use of all_reduce inconsistent when using world_size == 1 vs world_size > 1. The issue manifests, for example, when you are logging running average loss value:

epoch_loss = xm.all_reduce(xm.REDUCE_SUM, loss_accum, scale=1.0 / ((idx + 1) * world_size))

[ysiraichi]

I'm not sure whether this is expected behavior or not, but have you tried setting XLA_ALWAYS_ALLREDUCE=1?
@bhavya01

[SVISHNUVARDHAN3610]

Hi @afzalxo xm.all_reduce function is used to collect the value of the variable of epoch_loss from all nodes present in the environment
example: if you have 8 nodes then it will collect loss from 8 and average it with the scale function you given in the xm.all_reduce but when you have only 1 node it fails to apply the scale

example if world_size > 1
loss = fun(loss1 + loss2 +..........+ loss_n)