mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-25 19:21:53 +08:00
[misc] fix debugging code (#13487)
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
parent
4fb8142a0e
commit
7b203b7694
@ -94,20 +94,20 @@ pynccl.disabled = False
|
||||
s = torch.cuda.Stream()
|
||||
with torch.cuda.stream(s):
|
||||
data.fill_(1)
|
||||
pynccl.all_reduce(data, stream=s)
|
||||
value = data.mean().item()
|
||||
out = pynccl.all_reduce(data, stream=s)
|
||||
value = out.mean().item()
|
||||
assert value == world_size, f"Expected {world_size}, got {value}"
|
||||
|
||||
print("vLLM NCCL is successful!")
|
||||
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(cuda_graph=g, stream=s):
|
||||
pynccl.all_reduce(data, stream=torch.cuda.current_stream())
|
||||
out = pynccl.all_reduce(data, stream=torch.cuda.current_stream())
|
||||
|
||||
data.fill_(1)
|
||||
g.replay()
|
||||
torch.cuda.current_stream().synchronize()
|
||||
value = data.mean().item()
|
||||
value = out.mean().item()
|
||||
assert value == world_size, f"Expected {world_size}, got {value}"
|
||||
|
||||
print("vLLM NCCL with cuda graph is successful!")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user