c10d: add Collectives abstraction #125978

d4l3k · 2024-05-10T23:12:20Z

This adds a new Collectives API for doing distributed collectives operations. This is intended to replace the current Elastic store abstraction with more performant and debugable primitives.

Design doc: https://docs.google.com/document/d/147KcKJXEHvk1Q6tISLbJVvLejHg_1kIhBQeu-8RQxhY/edit

The standard implementation is using StoreCollectives but other more performant backends will be added in a follow up PR.

Test plan:

python test/distributed/test_collectives.py -v

This tests both functionality using multiple threads as well as timeout behavior.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang

pytorch-bot · 2024-05-10T23:12:24Z

🔗 Helpful Links

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

📄 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 e14d5f7 with merge base 697ed6f ():
💚 Looks good so far! There are no failures yet. 💚

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

shuqiangzhang · 2024-05-10T23:36:56Z

torch/csrc/distributed/c10d/Collectives.hpp

+class TORCH_API Collectives : public torch::CustomClassHolder {
+ public:
+  virtual void barrier(
+      const std::string& prefix,


Is 'prefix' specific to store based implementation of the APIs? If so, I imagine the APIs could be implemented with non-store semantics

It refers to the store based APIs but other implementations have a similar concept of an operation ID. I.e. the label field in https://fburl.com/code/dq8shbwi.

It may be best to change this name from prefix to "operation_uid"

I wonder if prefix better maps to e.g. sub_world_id in SC? Is the semantics that operations on each "prefix"/"subworld" must be sequential, and using different prefixes allows for a new namespace/fork to perform concurrent operations in e.g. separate threads?

But yeah, it does feel like prefix isn't a great name for the base class...

I think this is still a better fit for label -- all operations on store must be sequential regardless of prefix.

If you want a fully separate namespace you can use a different PrefixStore and create a new StoreCollectives instance -- similar to how you would use a different sub_world_id and create a new SCCollectives instance

Renamed to generic key

kwen2501 · 2024-05-11T06:57:59Z

Does the term Collectives here refer to "collectives for control plane (torchrun/elastic)"?
If so, can file name and class name like Collectives.hpp be renamed to more clearly differentiate from the data plane collectives (i.e. c10d)? Thanks!

kwen2501 · 2024-05-11T07:08:20Z

In addition, it seems torch/csrc/distributed/c10d has many files now. Does it make sense for store to have its own directory? I am just not sure whether the new directory should under c10d or at the same level of c10d. If store is used by both c10d and torchrun/elastic, it seems the latter makes sense. However, a lot of store files also have c10d namespace, so that makes me uncertain.

Skylion007

Performance nits / improvements.

torch/csrc/distributed/c10d/StoreCollectives.cpp

torch/csrc/distributed/c10d/init.cpp

torch/csrc/distributed/c10d/StoreCollectives.cpp

d4l3k · 2024-05-13T23:25:04Z

updates:

moved to c10d/control_collectives dir
renamed Collectives to ControlCollectives
renamed prefix to key to be backend agnostic
applies fmt/vector reserve optimizations

kurman

nice

kurman · 2024-05-13T23:30:28Z

torch/csrc/distributed/c10d/control_collectives/StoreCollectives.hpp

+
+namespace c10d {
+
+class TORCH_API StoreCollectives : public ControlCollectives {


StoreControlCollectives?

I thought about that but figured Store would be clear enough -- it wouldn't make sense to use the Store API for dataplane operations

kurman · 2024-05-13T23:31:43Z

torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp

+class TORCH_API ControlCollectives : public torch::CustomClassHolder {
+ public:
+  virtual void barrier(
+      const std::string& key,


key is better, I wonder if it is actually a 'token'?

I think they're equivalent -- key maps better to the underlying concept so will just leave it as is

torch/csrc/distributed/c10d/Store.hpp

shuqiangzhang · 2024-05-14T17:50:07Z

torch/csrc/distributed/c10d/control_collectives/StoreCollectives.cpp

+}
+
+void StoreCollectives::broadcast_send(
+    const std::string& key,


For these simple APIs such as broadcast, is 'key' really necessary for users? can we make it optional for users? e.g., they just need to provide the data to send and recv

@shuqiangzhang we need a unique identifier for every round due to some implementation details -- we could auto increment it internally but could cause problems. SC and store backends both require a unique key

shuqiangzhang · 2024-05-14T18:04:09Z

test/distributed/test_control_collectives.py

+
+        def f(rank: int) -> None:
+            collectives = dist.StoreCollectives(store, rank, world_size)
+            collectives.barrier("foo", timedelta(seconds=10), True)


Yeah, I mean this 'foo' does not really provide much value for users? could be optional or removed?

For store, this needs to be unique, if you call it multiple times with "foo" it won't actually do any barrier operation

OK, I was thinking it would be much nicer and cleaner if our general API could remove the burden from users of maintaining an unique key for each collective due to implementation limitations. Can we maintain a seqID or something within the implementation? I think this is also similar to our data plane side of collective APIs

we could potentially -- it gets a bit risky when you consider the case that not all workers may be joining at the same time. In elastic it's elastic so a simple seqID may not be sufficient and you need to key it off of the rendezvous round instead

OK, this makes sense in elastic, let's move forward

shuqiangzhang

LGTM, great start of the new control plane! Please do remember to unify the variable naming scheme

d4l3k · 2024-05-15T23:27:15Z

@pytorchbot merge

pytorchmergebot · 2024-05-15T23:29:49Z

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

pytorchmergebot · 2024-05-15T23:45:37Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-focal-cpu-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

pytorchmergebot · 2024-05-16T01:38:59Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

Lint / lintrunner-noclang / linux-job

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

d4l3k · 2024-05-17T00:30:22Z

@pytorchbot merge

pytorchmergebot · 2024-05-17T00:32:57Z

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

pytorchmergebot · 2024-05-17T01:10:06Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

Lint / lintrunner-noclang / linux-job

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

d4l3k · 2024-05-17T01:15:54Z

@pytorchbot merge

pytorchmergebot · 2024-05-17T01:18:50Z

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

This adds a new `Collectives` API for doing distributed collectives operations. This is intended to replace the [current Elastic store abstraction](https://github.com/pytorch/pytorch/blob/main/torch/distributed/elastic/utils/store.py) with more performant and debugable primitives. Design doc: https://docs.google.com/document/d/147KcKJXEHvk1Q6tISLbJVvLejHg_1kIhBQeu-8RQxhY/edit The standard implementation is using `StoreCollectives` but other more performant backends will be added in a follow up PR. Test plan: ``` python test/distributed/test_collectives.py -v ``` This tests both functionality using multiple threads as well as timeout behavior. Pull Request resolved: pytorch#125978 Approved by: https://github.com/shuqiangzhang

DanilBaibak · 2024-05-20T07:38:49Z

@pytorchbot revert -m "Break internal build" -c ghfirst

pytorchmergebot · 2024-05-20T07:40:36Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

This reverts commit 4b2ae2a. Reverted #125978 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](#125978 (comment)))

pytorchmergebot · 2024-05-20T07:40:45Z

@d4l3k your PR has been successfully reverted.

DanilBaibak · 2024-05-20T07:43:41Z

Hi @d4l3k! Sorry, I need to revert your PR because it breaks the internal build. Here you can find more information about the issue - D57518638.

This reverts commit d9c3485.

Summary: This reverts commit d9c3485. Reapplies #125978. cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang Reviewed By: c-p-i-o Differential Revision: D57572818 Pulled By: d4l3k

This reverts commit d9c3485. Reapplies #125978. Pull Request resolved: #126695 Approved by: https://github.com/c-p-i-o

d4l3k requested review from zdevito, suo, wconstab and kurman May 10, 2024 23:12

pytorch-bot bot added oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (c10d) release notes category labels May 10, 2024

d4l3k changed the title ~~c10d: added Collectives abstraction~~ c10d: add Collectives abstraction May 10, 2024

shuqiangzhang reviewed May 10, 2024

View reviewed changes

Skylion007 reviewed May 11, 2024

View reviewed changes

torch/csrc/distributed/c10d/StoreCollectives.cpp Outdated Show resolved Hide resolved

d4l3k force-pushed the d4l3k/collectives branch from b9e1a53 to 05adfea Compare May 13, 2024 23:24

d4l3k requested review from shuqiangzhang, Skylion007 and andrewjcg May 13, 2024 23:25

kurman reviewed May 13, 2024

View reviewed changes

shuqiangzhang reviewed May 14, 2024

View reviewed changes

torch/csrc/distributed/c10d/Store.hpp Outdated Show resolved Hide resolved

shuqiangzhang reviewed May 14, 2024

View reviewed changes

d4l3k force-pushed the d4l3k/collectives branch from 05adfea to 4a6661d Compare May 14, 2024 23:50

shuqiangzhang approved these changes May 15, 2024

View reviewed changes

d4l3k force-pushed the d4l3k/collectives branch from 4a6661d to f8e8100 Compare May 15, 2024 23:24

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 15, 2024

pytorchmergebot added the merging label May 15, 2024

pytorchmergebot removed the merging label May 16, 2024

d4l3k force-pushed the d4l3k/collectives branch from eefec66 to b345e70 Compare May 17, 2024 00:30

pytorchmergebot added the merging label May 17, 2024

pytorchmergebot removed the merging label May 17, 2024

c10d: added Collectives abstraction

e14d5f7

d4l3k force-pushed the d4l3k/collectives branch from b345e70 to e14d5f7 Compare May 17, 2024 01:15

pytorchmergebot added the merging label May 17, 2024

pytorchmergebot added the Merged label May 17, 2024

pytorchmergebot closed this in 4b2ae2a May 17, 2024

pytorchmergebot removed the merging label May 17, 2024

d4l3k deleted the d4l3k/collectives branch May 17, 2024 16:19

pytorchmergebot added a commit that referenced this pull request May 20, 2024

Revert "c10d: add Collectives abstraction (#125978)"

d9c3485

This reverts commit 4b2ae2a. Reverted #125978 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](#125978 (comment)))

pytorchmergebot added the Reverted label May 20, 2024

d4l3k added a commit that referenced this pull request May 20, 2024

Reapply "c10d: add Collectives abstraction (#125978)"

e98bda8

This reverts commit d9c3485.

d4l3k mentioned this pull request May 20, 2024

Reapply "c10d: add Collectives abstraction (#125978)" #126695

Closed

pytorchmergebot pushed a commit that referenced this pull request May 21, 2024

Reapply "c10d: add Collectives abstraction (#125978)" (#126695)

ac51920

This reverts commit d9c3485. Reapplies #125978. Pull Request resolved: #126695 Approved by: https://github.com/c-p-i-o


		namespace c10d {

		class TORCH_API StoreCollectives : public ControlCollectives {

c10d: add Collectives abstraction #125978

c10d: add Collectives abstraction #125978

Conversation

d4l3k commented May 10, 2024 • edited by pytorch-bot bot

pytorch-bot bot commented May 10, 2024 • edited

🔗 Helpful Links

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

✅ No Failures

Choose a reason for hiding this comment

d4l3k May 11, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kwen2501 commented May 11, 2024

kwen2501 commented May 11, 2024

Skylion007 left a comment

Choose a reason for hiding this comment

d4l3k commented May 13, 2024

kurman left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shuqiangzhang May 14, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shuqiangzhang May 14, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shuqiangzhang May 15, 2024 • edited

Choose a reason for hiding this comment

shuqiangzhang left a comment

Choose a reason for hiding this comment

d4l3k commented May 15, 2024

pytorchmergebot commented May 15, 2024

Merge started

pytorchmergebot commented May 15, 2024

Merge failed

pytorchmergebot commented May 16, 2024

Merge failed

d4l3k commented May 17, 2024

pytorchmergebot commented May 17, 2024

Merge started

pytorchmergebot commented May 17, 2024

Merge failed

d4l3k commented May 17, 2024

pytorchmergebot commented May 17, 2024

Merge started

DanilBaibak commented May 20, 2024

pytorchmergebot commented May 20, 2024

pytorchmergebot commented May 20, 2024

DanilBaibak commented May 20, 2024

d4l3k commented May 10, 2024 •

edited by pytorch-bot bot

pytorch-bot bot commented May 10, 2024 •

edited

d4l3k May 11, 2024 •

edited

shuqiangzhang May 14, 2024 •

edited

shuqiangzhang May 14, 2024 •

edited

shuqiangzhang May 15, 2024 •

edited