Use bucketized model averaging for LocalSGD

[d4l3k]

We want to implement the bucketization from https://github.com/pytorch/pytorch/blob/main/torch/distributed/algorithms/model_averaging/utils.py#L22-L50

Instead of our naive per parameter reduction here: https://github.com/pytorch/torchft/blob/main/torchft/local_sgd.py#L179-L181

[distributedstatemachine]

I would love to take this on .

[Krishn1412]

Hey @d4l3k , can I pick this up?

[d4l3k]

@Krishn1412 feel free to grab it

[Krishn1412]

@d4l3k can you take a look #111