Use device_id in dist init to reduce NCCL communicator warmup & creation overhead (#5728)

Edenzzzz · web-flow · commit dfb322642fe6 · 2025-04-26T18:11:09.000-07:00
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
@@ -1055,6 +1055,11 @@ def init_distributed_environment(
             world_size=world_size,
             rank=rank,
             timeout=timeout,
+            device_id=torch.device(
+                f"cuda:{torch.cuda.current_device()}"
+                if hasattr(torch, "cuda") and torch.cuda.is_available()
+                else None
+            ),  # Allow NCCL to eagerly init communicator
         )
 
     # set the local rank
