local_sgd: initial version of fault tolerant LocalSGD

[d4l3k]

This implements an initial version of LocalSGD. #39

This has some limitations:

• workers can only join at the start of the sync_every period
• in async mode, the joining worker will discard the first sync_every steps as it's treated as a single Manager step
• if any errors/workers fail we will only detect it at the end of sync_every steps
• This uses a copy of the weights in memory in order to restore the model on commit failure to guarantee we always step exactly sync_every steps on each worker. This copy will be required for DiLoCo to compute the pseudo gradient.
• We don't do any bucketization for the allreduce operation

Test plan:

pyre
pytest

This will require more testing as I've only tested it with a mocked Manager and no e2e/integration tests

[H-Huang]

in the function description it is saying we allreduce the weights, but the function implemented below allreduces the grads?

[d4l3k]

allreduce_grads is a bit misnamed, I'll rename it as we're not reducing gradients here

The behavior is just an allreduce though with some rescaling logic on top

[d4l3k]

done

[H-Huang]

hmm, i wonder instead of having the user call step after optimizer step, we modify the init of LocalSGD to pass in the optimizer as well and then add an optimizer hook (https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.register_step_post_hook.html#torch.optim.Optimizer.register_step_post_hook) to automatically call this logic?

[d4l3k]

TIL -- that looks pretty useful

I do think we can potentially move this into the forward wrapper as well but it makes things a bit harder to reason about

[d4l3k]

done

[H-Huang]

Very cool to add localSGD!