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

2024-11-25 14:55:01 -08:00
parent 4d62bca542
commit c4336b2b60
1 changed files with 2 additions and 1 deletions
--- 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"
            )