Skip to content

[detailed][benchmark] all_to_all_single with async_op #3436

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
Closed
wants to merge 1 commit into from
Closed

[detailed][benchmark] all_to_all_single with async_op #3436

wants to merge 1 commit into from
+118 −30

Conversation

TroyGarden
Copy link
Contributor

@TroyGarden TroyGarden commented Oct 5, 2025 •
edited
Loading

Summary:

context

  • add benchmark for all_to_all_single with async_op option.
  • the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
  • of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
  • when there's data dependency on the comms' output, user's need to call req.wait() explicitly, so that the main cuda stream wait on the comms stream

NOTE: the req.wait() call is non-blocking on the CPU side.

ref

optimization

  • the data validation is done on the device side, and very often the validation result is needed from the cpu side
  • in a previously example (a2a-sync), the assertion needs the checks on the cpu side, so it's blocking the cpu execution (see below)
image
  • actually the checks is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
image
  • we tried to use non_blocking=True in the copy_to, and the results are very interesting:
    all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to print(checks) to bypass the assertion in order to generate the trace)
image
  • the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (checks) becomes available on the cpu side.
image
  • as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
image

results

Differential Revision: D83924526

@meta-codesync Meta CodeSync
Copy link
Contributor

meta-codesync bot commented Oct 5, 2025

@TroyGarden has exported this pull request. If you are a Meta employee, you can view the originating Diff in D83924526.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 5, 2025
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 6, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
80525b3 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 6, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
b991118 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526
@TroyGarden TroyGarden changed the title all_to_all_single with async_op [benchmark] all_to_all_single with async_op Oct 6, 2025
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 6, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
fc28921 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 7, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
2e7f3aa 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Reviewed By: spmex

Differential Revision: D83924526
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 7, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
534c7e7 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 7, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
967605e 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526
@TroyGarden TroyGarden force-pushed the export-D83924526 branch 2 times, most recently from 967605e to a8177d1 Compare October 7, 2025 03:54
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 7, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
a8177d1 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Oct 7, 2025
@TroyGarden @facebook-github-bot
all_to_all_single with async_op (meta-pytorch#3436)
dd43f81 
Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526
@TroyGarden @meta-codesync
all_to_all_single with async_op (meta-pytorch#3436)
463971e 
Summary:
Pull Request resolved: meta-pytorch#3436

# context
* add benchmark for `all_to_all_single` with `async_op` option.
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184}
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242}
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281}
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334}
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765}

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526
@meta-codesync meta-codesync bot closed this in 4062615 Oct 7, 2025
@TroyGarden TroyGarden deleted the export-D83924526 branch October 7, 2025 05:23
@TroyGarden TroyGarden changed the title [benchmark] all_to_all_single with async_op [detailed][benchmark] all_to_all_single with async_op Oct 8, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. fb-exported meta-exported
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@TroyGarden