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

c4336b2b · Lianmin Zheng · GitHub · 4d62bca5 · c4336b2b
Unverified Commit c4336b2b authored Nov 25, 2024 by Lianmin Zheng Committed by GitHub Nov 25, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

python/sglang/srt/model_executor/cuda_graph_runner.py python/sglang/srt/model_executor/cuda_graph_runner.py +2 -1

No files found.
--- 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"
            )