Hello, and thanks for this clean implementation of Diloco.
It seems like loss_batch=0 at line 169 of pure_torch_diloco.py has incorrect indentation. The loss batch should be reset by all local ranks after the accumulated inner step is taken. Currently, it is only reset for rank 0; so loss_batch=0 should be outside the if block that starts at line 154.
Hello, and thanks for this clean implementation of Diloco.
It seems like
loss_batch=0at line 169 of pure_torch_diloco.py has incorrect indentation. The loss batch should be reset by all local ranks after the accumulated inner step is taken. Currently, it is only reset for rank 0; soloss_batch=0should be outside the if block that starts at line 154.