Fix allgather ops inside cuda graphs #3709
Open
+37
−7
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nvcastet
requested review from
merrymercy,
Ying1123,
hnyls2002,
zhyncs,
ispobock and
ByronHsu
as code owners
5 tasks
yizhang2077
reviewed
Feb 20, 2025
|
Seeing the same issue, please merge soon :) |
Hi @nvcastet May you help take a look? Thanks! |
|
Looking. |
|
Hi @nvcastet , I see cuda graph capture failed when use pynccl all gather in logits layer, but it will success in dp attention since |
|
@yizhang2077 Still debugging. Can I have another hour-ish? |
Ok, I think it is reasonable |
|
@yizhang2077 Is it reproducible on your side? |
It is weird, I will try it on my env. |
|
@nvcastet @yizhang2077 |
|
Hi @nvcastet, I can reproduce this error in my env. torch 2.5.1+cu124, 8xH200, docker image lmsysorg/sglang:dev |
|
@yizhang2077 Thanks let me try your container image.
From https://pytorch.org/docs/stable/notes/cuda.html#constraints |
|
The failure did not happen for me on pytorch 2.7 but it does on 2.5. which points to pynccl not being graphable (breaking the torch.compile graph) and also being recompiled later on during cuda graph capture. I added a commit to remove pynccl and replace it with a processgroup clone which should behave the same way. |
|
@nvcastet It is wierd, since pynccl allreduce is also in critical path and is graphable. |
|
Hi @nvcastet , pynccl cannot be removed since it's used for cu118 environment. PyTorch cu118 will install the nvidia-nccl-cu11 which with cuda 11.0 version of nccl by default. But cuda graph needs nccl with cuda version > 11.3. So we can specify the nccl so path |
The message (displayed using
@ispobock Thank you for pointing that out, I was not aware of this issue. |
|
@ispobock when downloading So we should be good, right? |
nvcastet
changed the title
|
Hi @nvcastet, I think currently the main problem is torch compile is conflicted with pynccl, right? I think you can try to wrap sglang/python/sglang/srt/distributed/parallel_state.py Lines 134 to 139 in 9087694 python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1. Could you take a try?
|
|
@nvcastet thx，it works！ |
|
I have patched and tested this pr but Watchdog caught collective operation timeout problems still have |
nvcastet
force-pushed
the
fix_all_gather_cuda_graph
branch
from
48a926f to
38585e0
Compare
|
@yizhang2077 @zhyncs I went back to the first commit and register the pynccl algather as a pytorch custom op as you suggested. |
Could you do me a favor? i have merge ur pr and the issue reproduce . Many thanks |
In my env, it will download the whl from |
|
@nvcastet Could you fix the lint with |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes #3424
TLDR:
sglang uses pynccl or their customalllreduce instead of pytorch ProcessGroup when graph capturing.
Allgather in sglang code base directly uses pytorch allgather instead of the sglang abstraction to decide which implementation to pick. Allgather is used in DP-attention and also to gather logits across TP dim.
The fix is to perform the allgather via the abstraction so that the same NCCL communicator won't be used inside and outside graph captures.
Motivation
Modifications
Checklist