fix

Author: hiworldwzj

lightllm/distributed/custom_all_gather.py

@@ -33,9 +33,6 @@
 
 if light_ops is not None:
     light_ops.meta_size()
-    custom_ar = True
-else:
-    custom_ar = False
 
 logger = init_logger(__name__)
 
@@ -65,7 +62,7 @@ def __init__(self, group: ProcessGroup, device: Union[int, str, torch.device], m
         self._IS_CAPTURING = False
         self.disabled = True
 
-        if not custom_ar:
+        if light_ops is None:
             # disable because of missing custom allgather library
             # e.g. in a non-cuda environment
             return

lightllm/distributed/custom_all_reduce.py

@@ -41,8 +41,8 @@
     # Use sgl custom allreduce
     ops = sgl_allreduce_ops
 
-ops.meta_size()
-custom_ar = True
+if ops is not None:
+    ops.meta_size()
 
 
 def is_weak_contiguous(inp: torch.Tensor):
@@ -70,7 +70,7 @@ def __init__(self, group: ProcessGroup, device: Union[int, str, torch.device], m
         self._IS_CAPTURING = False
         self.disabled = True
 
-        if not custom_ar:
+        if ops is None:
             # disable because of missing custom allreduce library
             # e.g. in a non-cuda environment
             return