Support current vLLM StatelessProcessGroup constructors

checkpoint_engine/distributed/vllm_compat.py

@@ -0,0 +1,16 @@
+import inspect
+from typing import Any
+
+
+def create_stateless_process_group(
+    group_cls: type,
+    *,
+    rank: int,
+    world_size: int,
+    store: Any,
+) -> Any:
+    """Construct a vLLM stateless process group across supported APIs."""
+    kwargs = {"rank": rank, "world_size": world_size, "store": store}
+    if "socket" in inspect.signature(group_cls).parameters:
+        kwargs["socket"] = None
+    return group_cls(**kwargs)

checkpoint_engine/distributed/vllm_hccl.py

@@ -19,6 +19,7 @@
 from vllm_ascend.utils import current_stream
 
 from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
 
 
 class HcclCommConfig(ctypes.Structure):
@@ -240,7 +241,12 @@ def init_process_group(
         self.world_size = world_size
         self.device = torch.device("npu", torch.npu.current_device())
 
-        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pg = create_stateless_process_group(
+            StatelessProcessGroup,
+            rank=rank,
+            world_size=world_size,
+            store=store,
+        )
         self.pyhccl = PyHcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pyhccl.comm
         self.initialized = True

checkpoint_engine/distributed/vllm_nccl.py

@@ -14,6 +14,7 @@
 from vllm.distributed.utils import StatelessProcessGroup
 
 from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
 
 
 try:
@@ -153,7 +154,12 @@ def init_process_group(
         self.world_size = world_size
         self.device = torch.device("cuda", torch.cuda.current_device())
 
-        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pg = create_stateless_process_group(
+            StatelessProcessGroup,
+            rank=rank,
+            world_size=world_size,
+            store=store,
+        )
         self.pynccl = PyNcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pynccl.comm
         self.initialized = True

tests/test_vllm_compat.py

@@ -0,0 +1,37 @@
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
+
+
+class CurrentProcessGroup:
+    def __init__(self, rank: int, world_size: int, store: object):
+        self.arguments = rank, world_size, store
+
+
+class LegacyProcessGroup:
+    def __init__(self, rank: int, world_size: int, store: object, socket: object):
+        self.arguments = rank, world_size, store, socket
+
+
+def test_current_vllm_process_group_signature():
+    store = object()
+
+    group = create_stateless_process_group(
+        CurrentProcessGroup,
+        rank=2,
+        world_size=8,
+        store=store,
+    )
+
+    assert group.arguments == (2, 8, store)
+
+
+def test_legacy_vllm_process_group_signature():
+    store = object()
+
+    group = create_stateless_process_group(
+        LegacyProcessGroup,
+        rank=2,
+        world_size=8,
+        store=store,
+    )
+
+    assert group.arguments == (2, 8, store, None)