Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

ProcessGroupBaby: support full suite of PG tests #89

Merged
merged 1 commit into from
Jan 29, 2025
Merged

ProcessGroupBaby: support full suite of PG tests #89

merged 1 commit into from
Jan 29, 2025

Conversation

d4l3k
Copy link
Member

@d4l3k d4l3k commented Jan 29, 2025

This makes ProcessGroupBaby pass the PG test suite and fixes some glaring gaps in functionality.

  • AllreduceOptions (and others) are now supported through a custom option wrapper
  • implements allgather, broadcast
  • BabyWork supports calling .wait() and .get_future() simultaneously
  • BabyWork will clean up even if .wait isn't called

This should allow ProcessGroupBabyNCCL to be used with TorchTitan

Test plan:

pytest torchft/process_group_test.py

@d4l3k d4l3k requested review from fegin and H-Huang January 29, 2025 22:50
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jan 29, 2025
fegin
fegin approved these changes Jan 29, 2025
View reviewed changes
Copy link
Contributor

@fegin fegin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

torchft/process_group.py
tx.put(op_id)
elif cmd == "del":
op_id: int = op[1]
del work[op_id]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we wait first before deleting it?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we need to? In normal situations the user code will have waited for it, or they will have gotten a future. If they have a future we still have a reference and don't need it in this map anymore

The only other common case would be if an error occured and we're cleaning things up/shutting down -- in that case waiting likely won't succeed so no reason to wait

@d4l3k d4l3k merged commit fa1630d into main Jan 29, 2025
6 checks passed
@d4l3k d4l3k deleted the d4l3k/baby branch January 29, 2025 23:53
@d4l3k d4l3k mentioned this pull request Jan 30, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@fegin fegin fegin approved these changes

@H-Huang H-Huang Awaiting requested review from H-Huang

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Meta Open Source bot.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@d4l3k @fegin @facebook-github-bot