implement collective all_to_all op #9442
Conversation
|
|
||
| return [t.cpu() for t in output_tensors] | ||
|
|
||
| def test_all_to_all(self): |
There was a problem hiding this comment.
do you think a performance test might be necessary here to ensure there's no unforseen bottleneck creating latency/throughput issues?
There was a problem hiding this comment.
Yes. Stage 1 of this project is to improve op coverage. Stage 2 is to rigorously benchmark how the collective ops scale and identify any bottlenecks. I've been reassigned to work on the new repo so I won't be doing stage 2, at least in the near future.
| @@ -359,6 +359,36 @@ def test_all_to_all_single(self, use_dynamo): | |||
| expected.sort().values), | |||
| f"Got {val}, expected {expected}") | |||
|
|
|||
| @staticmethod | |||
There was a problem hiding this comment.
are tests in pjrt/ designed to run on some non-trival distributed setup?
There was a problem hiding this comment.
In some cases yes. The tests in this file are run by tpu/run_tests.sh and expect multiple TPUs. Some of the other files in pjrt/ are part of the basic test suite (example)
How do you mean? There are two torch.distributed functions: |
#9315