[Distributed] all_reduce op and distributed info in graphs

examples/distributed/test.py

@@ -1,30 +1,32 @@
 """
 Testing script for distributed components for hidet
 To debug, set the environment variable NCCL_DEBUG=INFO
+
+To install nccl, run
+
+    pip install nvidia-nccl-cu11==2.18.3
+
+Or
+
+    pip install nvidia-nccl-cu12==2.18.3
 """
 import hidet
 import multiprocessing
 from multiprocessing import Process
 import numpy
 import argparse
 
+import hidet
 import hidet.cuda.nccl
 from hidet.cuda import nccl
-from hidet.cuda.nccl import NcclUniqueId, NcclDataType, NcclRedOp, nccl_library_filename
-from hidet.ffi import runtime_api
-from hidet.lang import attrs
-from hidet.ir.primitives.cuda.nccl import all_reduce
-from hidet.ir.type import data_type
-from hidet.utils import prod
-from hidet.drivers import build_ir_module
-from hidet.cuda.nccl.libinfo import get_nccl_include_dirs, get_nccl_library_search_dirs
-from hidet.runtime import load_compiled_module
+from hidet.cuda.nccl import NcclUniqueId
 
 print("NCCL version:", nccl.nccl_version())
 
 parser = argparse.ArgumentParser()
 parser.add_argument("n_gpus", type=int)
 parser.add_argument("reduce_op", choices=['sum', 'prod', 'max', 'min', 'avg'])
+parser.add_argument("--group_size", type=int, default=0)
 args = parser.parse_args()
 
 def run(world_size, rank, shared_id, barrier):
@@ -37,47 +39,43 @@ def run(world_size, rank, shared_id, barrier):
     barrier.wait()
     hidet.cuda.set_device(rank)
 
-    print('initialize', rank)
-    # Create NcclCommunicator and set the cuda context
-    # this part should be moved into CompiledGraph in the future
-    comm = nccl.create_comm(world_size, shared_id, rank)
-    comms_array = nccl.comms_to_array([comm])
-    runtime_api.set_nccl_comms(comms_array)
+    use_group = args.group_size > 1
+    if use_group:
+        gs = args.group_size
+        gn = world_size // gs
+        assert world_size % gs == 0
+        groups = [list(range(i * gs, (i + 1) * gs)) for i in range(gn)]
+    else:
+        groups = []
+
 
-    # Initialize send and receive buffer
     device = f"cuda:{rank}"
-    send = hidet.randn([2, 2], device=device)
-    recv = hidet.empty([2, 2], device=device)
-
-    print(rank, send)
-
-    dtype = data_type('float32')
-    shape = [2, 2] 
-    nbytes = dtype.nbytes * prod(shape)
-
-    # Define IRModule
-    with hidet.script_module() as script_module:
-        @hidet.script
-        def launch(send: dtype[shape], recv: dtype[shape]):
-            attrs.func_kind = 'public'
-            all_reduce(0, send, recv, nbytes, dtype, getattr(NcclRedOp, args.reduce_op))
-
-    # Build
-    ir_module = script_module.ir_module()
-    ir_module.target = 'cuda'
-    ir_module.include_dirs.extend(get_nccl_include_dirs())
-    ir_module.linking_dirs.extend(get_nccl_library_search_dirs())
-    ir_module.include_headers.append(["nccl.h"])
-    ir_module.linking_libs.append(":" + nccl_library_filename())
-    out_dir = f'./.cache/all_reduce_{rank}'
-
-    build_ir_module(ir_module, out_dir, target='cuda')
-    compiled_module = load_compiled_module(out_dir)
-
-    compiled_module(send, recv)
+    x = hidet.randn([1, 3], device=device)
+    w = hidet.randn([3, 2], device=device)
+
+    # Create Computation Graph
+    x_symb = hidet.symbol_like(x)
+    w_symb = hidet.symbol_like(w)
+    y_local = hidet.ops.relu(x_symb @ w_symb)
+    y_sync = hidet.ops.all_reduce(y_local, args.reduce_op, comm_id=int(use_group))
+    graph = hidet.trace_from([y_local, y_sync], inputs=[x_symb, w_symb])
+    opt_graph = hidet.graph.optimize(graph)
+    opt_graph.set_dist_attrs(nrank=world_size, rank=rank, groups=groups)
+    compiled = opt_graph.build()
+
+    # test save and load
+    compiled_dir = f"./outs/graph_{rank}.zip"
+    compiled.save(compiled_dir)
+    compiled = hidet.runtime.load_compiled_graph(compiled_dir)
+
+    # Create Distributed Graph
+    compiled.init_dist(shared_id)
+
+    y_local, y_sync = compiled(x, w)
+
     s = hidet.cuda.current_stream() 
     s.synchronize()
-    print(rank, recv)
+    print(f"process {rank}\nbefore allreduce:{y_local}\nafter allreduce:{y_sync}\n", end='')
 
 world_size = args.n_gpus
 

python/hidet/cuda/nccl/__init__.py

@@ -9,5 +9,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from .comm import create_comm, NcclUniqueId, NcclDataType, NcclRedOp, comms_to_array, init_unique_id, dtype_to_nccl
-from .ffi import nccl_version, nccl_library_filename
+from .ffi import nccl_available, nccl_version, nccl_library_filename
+from .comm import (
+    create_comm,
+    NcclUniqueId,
+    NcclDataType,
+    NcclRedOp,
+    comms_to_array,
+    init_unique_id,
+    dtype_to_nccl,
+    NcclCommunicator,
+    str_to_nccl_op,
+    NCCL_SPLIT_NOCOLOR,
+)

python/hidet/cuda/nccl/comm.py

@@ -15,7 +15,12 @@
 
 from hidet.ffi.utils import Array
 from hidet.ir.type import void_p, DataType
-from .ffi import nccl_runtime_api, NcclUniqueId
+from .ffi import nccl_available, NcclUniqueId
+
+NCCL_SPLIT_NOCOLOR = -1
+
+if nccl_available():
+    from .ffi import nccl_runtime_api
 
 
 class NcclDataType(IntEnum):
@@ -44,13 +49,20 @@ class NcclRedOp(IntEnum):
     avg = 4
 
 
+def str_to_nccl_op(name: str) -> NcclRedOp:
+    if name not in ('sum', 'prod', 'max', 'min', 'avg'):
+        raise RuntimeError(f"'{name}' is not a supported reduce op")
+    return getattr(NcclRedOp, name)
+
+
 class NcclCommunicator:
     def __init__(self, handle: int):
         """
         Users should not call this constructor directly. Because there are two ways of creating
         a new communicator: 1) using unique_id and rank ; 2) using split.
         """
-
+        if not nccl_available():
+            raise RuntimeError("NCCL is not available")
         self._handle = handle
 
     def __del__(self):
@@ -60,11 +72,16 @@ def __del__(self):
     def handle(self):
         return self._handle
 
-    def split(self):
-        raise NotImplementedError()
+    def split(self, key, color):
+        new_handle = nccl_runtime_api.comm_split(self._handle, color, key)
+        if color == NCCL_SPLIT_NOCOLOR:
+            return None
+        return NcclCommunicator(new_handle)
 
 
 def create_comm(nranks: int, unique_id: NcclUniqueId, rank: int) -> NcclCommunicator:
+    if not nccl_available():
+        raise RuntimeError("NCCL is not available")
     handle = nccl_runtime_api.comm_init_rank(nranks, unique_id, rank)
     return NcclCommunicator(handle)
 
@@ -77,6 +94,8 @@ def comms_to_array(comms: List[NcclCommunicator]) -> Array:
 
 
 def init_unique_id(unqie_id: NcclUniqueId) -> None:
+    if not nccl_available():
+        raise RuntimeError("NCCL is not available")
     nccl_runtime_api.get_unique_id(unqie_id)
 
 

python/hidet/cuda/nccl/ffi.py

@@ -49,8 +49,6 @@ def load_nccl_library():
         _LIB_NCCL = ctypes.cdll.LoadLibrary(lib_nccl_paths[0])
         nccl_library_path = lib_nccl_paths[0]
         break
-    if _LIB_NCCL is None:
-        raise OSError('Can not find nccl library in the following directory: \n' + '\n'.join(library_dirs))
 
 
 load_nccl_library()
@@ -60,49 +58,61 @@ def nccl_library_filename():
     return os.path.basename(nccl_library_path)
 
 
-if not nccl_available():
-    raise RuntimeError("NCCL Library not found.")
+if nccl_available():
 
-
-class NCCLRuntimeAPI:
-    """
-    Runtime APIs regarding NCCL
-    TODO: Exception handling
-    """
-
-    _get_version = get_func('ncclGetVersion', [c_void_p], c_int, lib=_LIB_NCCL)
-    _get_unique_id = get_func('ncclGetUniqueId', [c_void_p], c_int, lib=_LIB_NCCL)
-    _comm_init_rank = get_func('ncclCommInitRank', [c_void_p, c_int, NcclUniqueId, c_int], c_int, lib=_LIB_NCCL)
-    _comm_destroy = get_func('ncclCommDestroy', [c_void_p], c_int, lib=_LIB_NCCL)
-
-    _comm_user_rank = get_func('ncclCommUserRank', [c_void_p, POINTER(c_int)], c_int, lib=_LIB_NCCL)
-    _comm_count = get_func('ncclCommCount', [c_void_p, POINTER(c_int)], c_int, lib=_LIB_NCCL)
-
-    @staticmethod
-    def get_version() -> int:
-        version = c_int(0)
-        NCCLRuntimeAPI._get_version(pointer(version))
-        return version.value
-
-    @staticmethod
-    def get_unique_id(comm_id: NcclUniqueId) -> None:
+    class NCCLRuntimeAPI:
         """
-        In-place initialization of the NcclUniqueId object
+        Runtime APIs regarding NCCL
+        TODO: Exception handling
         """
-        ret = NCCLRuntimeAPI._get_unique_id(pointer(comm_id))
-        assert ret == 0, ret
-
-    @staticmethod
-    def comm_init_rank(ndev: int, comm_id: NcclUniqueId, rank: int) -> int:
-        comm = c_void_p()
-        ret = NCCLRuntimeAPI._comm_init_rank(pointer(comm), ndev, comm_id, rank)
-        assert ret == 0, ret
-        return comm.value
-
-    @staticmethod
-    def comm_destroy(comm_handle) -> None:
-        ret = NCCLRuntimeAPI._comm_destroy(comm_handle)
-        assert ret == 0
-
 
-nccl_runtime_api = NCCLRuntimeAPI()
+        _get_version = get_func('ncclGetVersion', [c_void_p], c_int, lib=_LIB_NCCL)
+        _get_unique_id = get_func('ncclGetUniqueId', [c_void_p], c_int, lib=_LIB_NCCL)
+        _comm_init_rank = get_func('ncclCommInitRank', [c_void_p, c_int, NcclUniqueId, c_int], c_int, lib=_LIB_NCCL)
+        _comm_destroy = get_func('ncclCommDestroy', [c_void_p], c_int, lib=_LIB_NCCL)
+
+        _comm_user_rank = get_func('ncclCommUserRank', [c_void_p, POINTER(c_int)], c_int, lib=_LIB_NCCL)
+        _comm_count = get_func('ncclCommCount', [c_void_p, POINTER(c_int)], c_int, lib=_LIB_NCCL)
+
+        # Early versions of NCCL do not have split
+        try:
+            _comm_split = get_func('ncclCommSplit', [c_void_p, c_int, c_int, c_void_p, c_void_p], c_int, lib=_LIB_NCCL)
+        except ValueError:
+            _comm_split = None
+
+        @staticmethod
+        def get_version() -> int:
+            version = c_int(0)
+            NCCLRuntimeAPI._get_version(pointer(version))
+            return version.value
+
+        @staticmethod
+        def get_unique_id(comm_id: NcclUniqueId) -> None:
+            """
+            In-place initialization of the NcclUniqueId object
+            """
+            ret = NCCLRuntimeAPI._get_unique_id(pointer(comm_id))
+            assert ret == 0, ret
+
+        @staticmethod
+        def comm_init_rank(ndev: int, comm_id: NcclUniqueId, rank: int) -> int:
+            comm = c_void_p()
+            ret = NCCLRuntimeAPI._comm_init_rank(pointer(comm), ndev, comm_id, rank)
+            assert ret == 0, ret
+            return comm.value
+
+        @staticmethod
+        def comm_destroy(comm_handle) -> None:
+            ret = NCCLRuntimeAPI._comm_destroy(comm_handle)
+            assert ret == 0
+
+        @staticmethod
+        def comm_split(comm_handle: int, color: int, key: int) -> int:
+            if NCCLRuntimeAPI._comm_split is None:
+                raise RuntimeError("split is not supported on this version of NCCL. Please install a newer version.")
+            comm = c_void_p()
+            ret = NCCLRuntimeAPI._comm_split(comm_handle, color, key, pointer(comm), None)
+            assert ret == 0
+            return comm.value
+
+    nccl_runtime_api = NCCLRuntimeAPI()

python/hidet/drivers/build_graph.py

@@ -22,7 +22,13 @@
 from hidet.graph.tensor import Tensor
 from hidet.graph.flow_graph import FlowGraph
 from hidet.runtime.compiled_module import CompiledModule
-from hidet.runtime.compiled_graph import CompiledGraph, GraphMetaData, GraphExecution, GraphExecutionInstruction
+from hidet.runtime.compiled_graph import (
+    CompiledGraph,
+    GraphMetaData,
+    GraphExecution,
+    GraphExecutionInstruction,
+    GraphDistributedInfo,
+)
 from hidet.runtime.compiled_task import CompiledTask, TensorSignature
 from hidet.graph.operator import Operator
 from hidet.ir import primitives
@@ -142,6 +148,12 @@ def get_graph_meta_data(graph: FlowGraph, num_kernels, space: int) -> GraphMetaD
     )
 
 
+def get_graph_dist_info(graph: FlowGraph) -> GraphDistributedInfo:
+    if not graph.is_distributed():
+        return None
+    return GraphDistributedInfo(nrank=graph.nrank, rank=graph.rank, groups=graph.groups)
+
+
 def build_graph_module(graph: FlowGraph, graph_weights: List[Tensor], node2kernel: List[int]) -> CompiledModule:
     from hidet.lang import void_p, attrs, int32, int64, meta, cast
     from hidet.ir.primitives.runtime import memory_planner_init, memory_planner_allocate, memory_planner_free
@@ -329,6 +341,9 @@ def build_flow_graph(graph, *, space=0) -> CompiledGraph:
     # get the graph meta data
     graph_meta_data = get_graph_meta_data(graph, len(graph_kernels), space)
 
+    # get distributed information
+    graph_dist_info = get_graph_dist_info(graph)
+
     # build the compiled graph
     compiled_graph = CompiledGraph(
         meta=graph_meta_data,
@@ -337,6 +352,7 @@ def build_flow_graph(graph, *, space=0) -> CompiledGraph:
         compiled_tasks=graph_kernels,
         graph_execution=graph_execution,
         graph_string=str(graph),
+        dist_info=graph_dist_info,
     )
 
     # save the compiled graph to cache