ctype load nccl pypi

a.py

@@ -0,0 +1,6 @@
+def a():
+    print(b())
+a()
+def b():
+    return 1
+

bmtrain/init.py

@@ -6,8 +6,15 @@
 from .utils import print_dict
 import ctypes
 from .global_var import config
-from . import nccl
+
+try:
+    from . import nccl
+except:
+    from .utils import load_nccl_pypi
+    load_nccl_pypi()
 from .synchronize import synchronize
+
+
 def init_distributed(
         init_method : str = "env://",
         seed : int = 0,
@@ -52,16 +59,9 @@ def init_distributed(
     port = os.environ["MASTER_PORT"]
     master = addr+":"+port
     timeout = datetime.timedelta(seconds=1800)
-    try:
-        rendezvous_iterator = dist.rendezvous(
-            init_method, rank, world_size, timeout=timeout
-        )   
-    except RuntimeError:
-        import nvidia.nccl
-        path = os.path.join(os.path.dirname(nvidia.nccl.__file__), "lib")
-        for file_so in os.listdir(path):
-            if file_so.endswith(".so"):
-                ctypes.CDLL(os.path.join(path, file_so))
+    rendezvous_iterator = dist.rendezvous(
+        init_method, rank, world_size, timeout=timeout
+    )   
 
     store, rank, world_size = next(rendezvous_iterator)
     store.set_timeout(timeout)
@@ -173,3 +173,4 @@ def get_group_rank(self,group_name):
 
 def is_initialized() -> bool:
     return config["initialized"]
+

bmtrain/utils.py

@@ -2,6 +2,8 @@
 import sys
 from typing import Any, Dict, Iterable, Optional
 from .global_var import config
+import os
+import ctypes
 
 ALIGN = 4
 ROW_WIDTH = 60
@@ -18,6 +20,17 @@ def check_torch_version(version_str):
     current_version_int = current_version_int_arr[0] * 10000 + current_version_int_arr[1] * 100 + current_version_int_arr[2]
     return current_version_int - version_int
 
+def load_nccl_pypi():
+    try:
+        import nvidia.nccl
+    except:
+        print("Run pip install nvidia-nccl-cu11 >=2.14.3 first")
+    path = os.path.join(os.path.dirname(nvidia.nccl.__file__), "lib")
+    for file_so in os.listdir(path):
+        if file_so.endswith(".so"):
+            ctypes.CDLL(os.path.join(path, file_so))
+
+
 def round_up(x, d):
     return (x + d - 1) // d * d