Add fake process group

[ezyang]

Stack from ghstack (oldest at bottom):

• -> Add fake process group #102180

Signed-off-by: Edward Z. Yang ezyang@meta.com

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/102180

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ac40616:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[wanchaol]

looks great, thanks!

[wanchaol]

we probably don't need a FakeStore and instead we can just use HashStore I suppose. But that could be in a separate PR.

[ezyang]

Can't use HashStore. For example, FSDP will attempt to do a barrier. The barrier will block you until enough writes into the store have happened. If we're doing fake PG there will be no other writes and you'll deadlock. It's best to have the store error if you try to do anything with it and route around it differently.

[wanchaol]

Hmmm I think FakePG barrier should be a no-op, and in terms of init_processs_group store based barrier it seems we already skip the barrier so it won't write anything to the HashStore. So either fake_store or hash_store could work I feel (I can give it a try and see if that's feasible or not)

[ezyang]

FSDP does a barrier which is why I ended up doing FakeStore. But yeah, try some stuff out, the goal is to be able to run FSDP end-to-end with the fake group with only one node.

[ezyang]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[albanD]

Very cool!
Do you have a small script showing the right way to use this?

[ezyang]

the test cases are pretty good

[albanD]

Ok!
But can I actually forward/backward with that Module? What should I expect the memory behavior of this thing to be (representative of what would happen with nccl or not?)

[wanchaol]

Ok! But can I actually forward/backward with that Module? What should I expect the memory behavior of this thing to be (representative of what would happen with nccl or not?)

@albanD the memory behavior should be representative of each rank behavior as if you are running a real multiprocessing job with nccl process group. i.e. if specify rank=1 and torch.cuda.set_device(1) when you initialize the fake pg, all of the memory behaviors should be similar as if you are in rank 1 of a real multiprocessing job. (as under the hood each rank still allocate the same amount of data for c10d collectivese)

[insujang]

May I ask why you implemented additional FakeProcessGroup implementation while we have MockProcessGroup in torch/testing/_internal/distributed/distributed_utils.py? I understand that MockProcessGroup is just mocking a process group and does nothing while FakeProcessGroup aims to provide fake communication as well, but it is a superset and seems no reason to maintain MockProcessGroup. Will you deprecate MockProcessGroup later?