Different nccl stream within EmbeddingKey alltoall and EmbeddingValue…

[zhaozheng09]

Now all nccl kernels are launched in one stream in pipeline mode, all nccl kernels will be waited each other if these ops conflict. We classified all nccl ops:

1. Alltoall embedding keys in next step.
2. Others nccl ops(like backward all2all, etc) in current step.
   They have no interdependence at all, so I change to a new stream .

Compare:
before:

after:

[facebook-github-bot]

Hi @zhaozheng09!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[dstaay-fb]

It's not recommended to create a separate Process Group here; and could result in performance regressions / reduced reliability. This guidance is from prior studies; may no longer be valid, but would require benchmarking/understanding as to why.

[zhaozheng09]

Indeed, when there are too many tables, there will be a lot of overhead in stream synchronization .
benchmarking: I will try put a small example .
why: alltoall embedding keys is independent with other nccl communication, it not need wait after alltoall sparse gradient and allreduce dense gradient in previous back propagation. if you put alltoall embedding key and alltoall sparse gradient and allreduce dense gradient to a stream, they maybe be block with each other in pipeline mode.

[dstaay-fb]

As per trace, it seems your able to trigger multiple comms stream, but at least based on my historical understanding this divides bandwidth of nccl stream; and results in less total throughput.

Worth maybe sharing results + confirming cpu is running ahead of gpu for clean comparison.

[zhaozheng09]

I reproduced this PR on DLRM. I got 5% preformance.

1. example: https://github.com/zhaozheng09/dlrm/tree/multiple_stream_torchrec
2. machine: A100
3. cmd:

1. 1stream(baseline): cd dlrm/torchrec_dlrm/ && torchx run -s local_cwd dist.ddp -j 1x2 --script dlrm_main.py -- --batch_size 8192 2>1stream.txt
2. 2streams(add this PR)cd dlrm/torchrec_dlrm/ && torchx run -s local_cwd dist.ddp -j 1x2 --script dlrm_main.py -- --batch_size 8192 2>2streams.txt
3. fgrep 'it/s' 1stream.txt | tail -n 150 | awk -F', |it/s' '{sum+=$2}END{print sum/NR}'
4. fgrep 'it/s' 2streams.txt | tail -n 150 | awk -F', |it/s' '{sum+=$2}END{print sum/NR}'
5. cmp 3,4 result .

4. version:
   torch: 2.4.1+cu121
   torchrec: 1.0.0
   fbgemm_gpu: 2024.11.25
   cuda: 12.2

5. result:
   1stream: 109.102it/s
   2streams: 114.948it/s(↑5.35%)

6. profiling:

1stream:

2streams:

@dstaay-fb

[zhaozheng09]

What should I do ? or What else do I need to prove? or Do you approve of this PR ? If you approve this PR, I need to close it quickly. thx. @dstaay-fb