Use c10d broadcast_object in Zero

[kumpera]

🚀 The feature, motivation and pitch

The current implementation of Zero Redundancy optimizer has its own implementation of object broadcasting.

We should replace it with c10d broadcast_object_list.

To verify the result, run the tests in test/distributed/optim/test_zero_redundancy_optimizer.py.

Alternatives

No response

Additional context

No response

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501

[awgu]

Does broadcast_object_list() still load Tensors in the broadcasted objects to the sender's GPU, or has that been changed?

This was an issue previously, so I tried landing #61539 to add a map_location argument similar to torch.load() to specify which device to load the Tensors in the broadcasted objects. However, since I changed the underlying serialization to use torch.save() / torch.load() instead of pickle.Pickler / pickle.Unpickler, it broke internal workloads and was reverted, where the hypothesis for the breakage was the difference in default pickle protocol level between pickle.Pickle and torch.save().

The ZeRO _broadcast_object() uses torch.save() / torch.load() and uses the map_location argument to avoid loading received Tensors to the sender's GPU.

[DanielAtt2000]

Heyy, New to contributing in open source. I'll give it a go but any help is much appreciated!

[harxish]

@DanielAtt2000 Hey, are you working on this?

[jazzysoggy]

I made an attempt at this, was wondering if anyone could review