Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "517858adf5602db78d269e1f8317b9db307ac0b4"
Unverified Commit c4336b2b authored by Lianmin Zheng's avatar Lianmin Zheng Committed by GitHub
Browse files

Use custom allreduce w/ torch.compile (#2185)

parent 4d62bca5
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
python/sglang/srt/model_executor/cuda_graph_runner.py
View file @ c4336b2b
...@@ -65,7 +65,8 @@ def patch_model( ...@@ -65,7 +65,8 @@ def patch_model(
_to_torch(model) _to_torch(model)
monkey_patch_vllm_all_gather() monkey_patch_vllm_all_gather()
backup_ca_comm = tp_group.ca_comm backup_ca_comm = tp_group.ca_comm
tp_group.ca_comm = None # Use custom-allreduce here
# tp_group.ca_comm = None
yield torch.compile( yield torch.compile(
torch.no_grad()(model.forward), mode="max-autotune-no-cudagraphs" torch.no_grad()(model.forward), mode="max-autotune-no-cudagraphs"
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment