From c4336b2b60acdc2a835842f5033c05226d211e56 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 25 Nov 2024 14:55:01 -0800 Subject: [PATCH] Use custom allreduce w/ torch.compile (#2185) --- python/sglang/srt/model_executor/cuda_graph_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 02bd358b0..596fd2104 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -65,7 +65,8 @@ def patch_model( _to_torch(model) monkey_patch_vllm_all_gather() backup_ca_comm = tp_group.ca_comm - tp_group.ca_comm = None + # Use custom-allreduce here + # tp_group.ca_comm = None yield torch.compile( torch.no_grad()(model.forward), mode="max-autotune-no-cudagraphs" )