[RFC] Do we really need all_reduce in `BaseLossContext`?

[nil0x9]

Currently in loss calculation, there is an invocation of all-reduce with autograd:

xtuner/xtuner/v1/loss/base_loss_ctx.py

Line 156 in 7b1246a

loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)

IIUC this all-reduce is primarily for logging purposes. It could be potentially be removed safely (plz correct me if I'm wrong here), bc:

1. FSDP works fine with just local loss, and
2. with the above reduced loss, during backward stage the gradient would be identical as local_loss;
3. loss would reduced in TrainEngine (thus it would be reduced twice!)
   

   xtuner/xtuner/v1/engine/train_engine.py

   Line 337 in 7b1246a

   dist.all_reduce(reduced_llm_loss.div_(dist.get_world_size()))

[HAOCHENYE]

I looked at

xtuner/xtuner/v1/loss/base_loss_ctx.py

Line 156 in 7b1246a

loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)

- this uses allreduce sum, while

xtuner/xtuner/v1/engine/train_engine.py

Line 337 in 7b1246a

dist.all_reduce(reduced_llm_loss.div_(dist.get_world_size()))

uses the equivalent allreduce mean.

Of course, the latter operation is indeed redundant (For CE Loss).

One reason the former needs to do this is for loss calibration - the loss on each rank is actually a very small value that has been normalized based on the total tokens consumed globally in the current step, and also takes into account factors like gradient accumulation, SP (sequence parallelism), etc.

Therefore, the loss obtained before reduce sum is actually a very, very small value. The reduce sum here restores it to a scale suitable for gradient updates (but it's still smaller by a factor of gradient accumulation steps).

Should we just remove the latter allreduce mean? Maybe not, since we need to maintain the principle of minimal knowledge in TrainEngine - it shouldn't need to know whether the llm_loss returned by the model has already been reduced or not. It could be safer to apply a reduce mean if the cost is negligible.

[nil0x9]

The reduce sum here restores it to a scale suitable for gradient updates

I don't get it. The gradient for all-reduce sum w.r.t. local param would be exactly the same as local loss, no? The gradient of $\sum_i x_i$ w.r.t. $x_i$ is identity ? So IIUC the scale of gradient is already suitable.