mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-07-16 16:07:18 +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()
|
s = torch.cuda.Stream()
|
||||||
with torch.cuda.stream(s):
|
with torch.cuda.stream(s):
|
||||||
data.fill_(1)
|
data.fill_(1)
|
||||||
pynccl.all_reduce(data, stream=s)
|
out = pynccl.all_reduce(data, stream=s)
|
||||||
value = data.mean().item()
|
value = out.mean().item()
|
||||||
assert value == world_size, f"Expected {world_size}, got {value}"
|
assert value == world_size, f"Expected {world_size}, got {value}"
|
||||||
|
|
||||||
print("vLLM NCCL is successful!")
|
print("vLLM NCCL is successful!")
|
||||||
|
|
||||||
g = torch.cuda.CUDAGraph()
|
g = torch.cuda.CUDAGraph()
|
||||||
with torch.cuda.graph(cuda_graph=g, stream=s):
|
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)
|
data.fill_(1)
|
||||||
g.replay()
|
g.replay()
|
||||||
torch.cuda.current_stream().synchronize()
|
torch.cuda.current_stream().synchronize()
|
||||||
value = data.mean().item()
|
value = out.mean().item()
|
||||||
assert value == world_size, f"Expected {world_size}, got {value}"
|
assert value == world_size, f"Expected {world_size}, got {value}"
|
||||||
|
|
||||||
print("vLLM NCCL with cuda graph is successful!")
|
print("vLLM NCCL with cuda graph is successful!")
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user