From 1d55c27f1e843963545d218730a4a1f048886f98 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Wed, 13 Sep 2023 23:04:07 -0700 Subject: [PATCH 1/5] unify ccl.py for cpu and gpu' --- csrc/cpu/comm/ccl.cpp | 20 ++++++-- deepspeed/comm/ccl.py | 106 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 113 insertions(+), 13 deletions(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 7fbf5f61392e..0c8b187104f1 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -278,9 +278,19 @@ int world_size = -1; std::set _comm_ids; std::set _colors; ccl::vector_class _ccl_comms; +std::map, int> group_to_comm_id; ccl::communicator& _get_comm_from_group() { return _ccl_comms[0]; } ccl::communicator& _get_comm_from_group(py::object group) { return _ccl_comms[0]; } +ccl::communicator& _get_comm_from_group(std::vector ranks) +{ + if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) + { + auto id = group_to_comm_id.find(ranks); + return _ccl_comms[id->second]; + } + return _ccl_comms[0]; +} #define CCLCHECK(cmd) \ do { \ @@ -452,7 +462,7 @@ ccl::reduction get_ccl_reduce_op(py::object op, at::Tensor& input) return ccl_op; } -void broadcast(torch::Tensor& data, int src, py::object group, bool async_op) +void broadcast(torch::Tensor& data, int src, std::vector group, bool async_op) { CCLCHECK(ccl::broadcast(data.data_ptr(), data.numel(), @@ -463,7 +473,7 @@ void broadcast(torch::Tensor& data, int src, py::object group, bool async_op) } // TODO: implement torch's async_op behavior, document it. -void all_reduce(torch::Tensor& data, py::object op, py::object group, bool async_op) +void all_reduce(torch::Tensor& data, py::object op, std::vector group, bool async_op) { CCLCHECK(ccl::allreduce(data.data_ptr(), data.data_ptr(), @@ -477,7 +487,7 @@ void all_reduce(torch::Tensor& data, py::object op, py::object group, bool async void all_reduce_caching(torch::Tensor& data, py::object op, std::string match_id, - py::object group, + std::vector group, bool async_op) { ccl::allreduce_attr attr = ccl::default_allreduce_attr; @@ -499,7 +509,7 @@ void all_reduce_caching(torch::Tensor& data, .wait()); } -void inference_all_reduce(torch::Tensor& data, py::object op, py::object group, bool async_op) +void inference_all_reduce(torch::Tensor& data, py::object op, std::vector group, bool async_op) { static py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); static auto ReduceOpSum = (int)py::int_(ReduceOp.attr("SUM").attr("value")); @@ -569,7 +579,7 @@ void inference_all_reduce(torch::Tensor& data, py::object op, py::object group, } } -void barrier(py::object group, bool async_op) +void barrier(std::vector group, bool async_op) { CCLCHECK(ccl::barrier(_get_comm_from_group(group)).wait()); } diff --git a/deepspeed/comm/ccl.py b/deepspeed/comm/ccl.py index 65646f6ba646..079df64c8cb8 100644 --- a/deepspeed/comm/ccl.py +++ b/deepspeed/comm/ccl.py @@ -20,6 +20,13 @@ def build_ccl_op(): print(f'DeepSpeed {builder.absolute_name()} built successfully') return ccl_cpp_module +class CCLHandler(): + def __init__(self, ccl_comm_op=None): + self.ccl_comm_op = ccl_comm_op + + def wait(self): + # backend covered it + pass class CCLBackend(TorchBackend): @@ -39,7 +46,7 @@ def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, i size = self.get_world_size() rank = self.get_rank() main_kvs = self.ccl_comm_op.get_kvs_addr(rank) - main_kvs = torch.tensor(main_kvs).to(torch.uint8) + main_kvs = torch.tensor(main_kvs).to(torch.uint8).to(get_accelerator().device_name(rank)) super(CCLBackend, self).broadcast(main_kvs, 0) self.ccl_comm_op.initialize(size, rank, main_kvs) self.initialized = True @@ -47,19 +54,102 @@ def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, i def is_initialized(self): return self.initialized - def broadcast(self, tensor, src, group=None, async_op=False): - self.ccl_comm_op.broadcast(tensor, src, group, async_op) - def all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): use_caching = False + group_ranks = self.get_all_ranks_from_group(group) if use_caching: match_id = f"{tensor.size()}-{op}" - self.ccl_comm_op.all_reduce_caching(tensor, op, match_id, group, async_op) + self.ccl_comm_op.all_reduce_caching(tensor, op, match_id, group_ranks, async_op) else: - self.ccl_comm_op.all_reduce(tensor, op, group, async_op) + self.ccl_comm_op.all_reduce(tensor, op, group_ranks, async_op) def inference_all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): self.ccl_comm_op.inference_all_reduce(tensor, op, group, async_op) - def barrier(self, group=None, async_op=False): - self.ccl_comm_op.barrier(group, async_op) + def broadcast(self, tensor, src, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.broadcast(tensor, src, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def all_gather(self, tensor_list, tensor, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.all_gather(tensor_list, tensor, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def reduce_scatter_tensor(self, output_tensor,input_tensor, op, group=None, async_op=False): + #todo: ccl version + super(CCLBackend, self).reduce_scatter_tensor(output_tensor,input_tensor, op, group) + + def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False): + #todo: ccl version + super(CCLBackend, self).all_gather_into_tensor(output_tensor, input_tensor, group) + + def all_to_all_single(self, output, input, output_split_sizes, input_split_sizes, group=None, async_op=False): + #todo: ccl version + super(CCLBackend, self).all_to_all_single(output, input, output_split_sizes, input_split_sizes, group) + + def send(self, tensor, dst, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.send(tensor, dst, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def recv(self, tensor, src, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.recv(tensor, src, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def gather(self, tensor, gather_list, dst, group=None, async_op=False): + #todo: ccl version + super(CCLBackend, self).gather(tensor, gather_list, dst, group) + + def scatter(self, tensor, gather_list, dst, group=None, async_op=False): + #todo: ccl version + super(CCLBackend, self).scatter(tensor, gather_list, dst, group) + + def barrier(self, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.barrier(group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def monitored_barrier(self, group=None, timeout=None, wait_all_ranks=False): + #todo: ccl version + super(CCLBackend, self).monitored_barrier(group) + + def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.reduce_scatter(output, input_list, op, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): + group_ranks = self.get_all_ranks_from_group(group) + self.ccl_comm_op.reduce(tensor, dst, op, group_ranks, async_op) + return CCLHandler(self.ccl_comm_op) + + def new_group(self, ranks): + size = len(ranks) + rank = self.get_rank() + if tuple(ranks) in self.groups or rank not in ranks: + return + sub_main_kvs = self.ccl_comm_op.get_sub_kvs_addr(rank == ranks[0]) + sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to("xpu:"+str(rank)) + torch_new_group = super(CCLBackend, self).new_group(ranks) + super(CCLBackend, self).broadcast(sub_main_kvs, ranks[0], torch_new_group, False) + self.ccl_comm_op.initialize_sub_comm(size, ranks.index(rank), sub_main_kvs, ranks) + self.groups.append(tuple(ranks)) + return torch_new_group + + def get_all_ranks_from_group(self, group): + if group is None: + return list(range(self.get_world_size())) + rank=0 + results=[] + try: + while True: + results.append(torch.distributed.distributed_c10d._get_global_rank(group, rank)) + rank+=1 + except RuntimeError: + pass + + if tuple(results) not in self.groups: + self.new_group(results) + return results From fe27d7675b3e8ec373350ebca94c5740a5c45f26 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Wed, 20 Sep 2023 20:27:50 -0700 Subject: [PATCH 2/5] handle all coll in on func --- csrc/cpu/comm/ccl.cpp | 7 +++++ deepspeed/comm/ccl.py | 68 ++++++++++++++++++------------------------- 2 files changed, 36 insertions(+), 39 deletions(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 0c8b187104f1..3291e74fc9c4 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -584,6 +584,12 @@ void barrier(std::vector group, bool async_op) CCLCHECK(ccl::barrier(_get_comm_from_group(group)).wait()); } +std::vector get_available_coll() +{ + std::vector colls{"broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; + return colls; +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("get_kvs_addr", &get_kvs_addr, "create and get main kvs addr"); @@ -595,4 +601,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("inference_all_reduce", &inference_all_reduce, "low latency all_reduce implementation"); m.def("all_reduce_caching", &all_reduce_caching, "ccl all_reduce with caching"); m.def("barrier", &barrier, "barrier"); + m.def("get_available_coll", &get_available_coll, "get_available_coll"); } diff --git a/deepspeed/comm/ccl.py b/deepspeed/comm/ccl.py index 079df64c8cb8..eef7e6c5a37b 100644 --- a/deepspeed/comm/ccl.py +++ b/deepspeed/comm/ccl.py @@ -50,80 +50,70 @@ def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, i super(CCLBackend, self).broadcast(main_kvs, 0) self.ccl_comm_op.initialize(size, rank, main_kvs) self.initialized = True + self.available_coll = self.ccl_comm_op.get_available_coll() def is_initialized(self): return self.initialized - + + def run_collective(self, name, **kwargs): + if name in self.available_coll: + kwargs['group'] = self.get_all_ranks_from_group(kwargs['group']) + func = "self.ccl_comm_op." + name + eval(func)(*(kwargs.values())) + return CCLHandler(self.ccl_comm_op) + else: + func = "super(CCLBackend, self)." + name + return eval(func)(*(kwargs.values())) + def all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): use_caching = False - group_ranks = self.get_all_ranks_from_group(group) if use_caching: match_id = f"{tensor.size()}-{op}" - self.ccl_comm_op.all_reduce_caching(tensor, op, match_id, group_ranks, async_op) + return self.run_collective(name="all_reduce_caching", tensor=tensor, op=op, match_id=match_id, group=group, async_op=async_op) else: - self.ccl_comm_op.all_reduce(tensor, op, group_ranks, async_op) + return self.run_collective(name="all_reduce", tensor=tensor, op=op, group=group, async_op=async_op) def inference_all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): - self.ccl_comm_op.inference_all_reduce(tensor, op, group, async_op) + return self.run_collective(name="inference_all_reduce", tensor=tensor, op=op, group=group, async_op=async_op) def broadcast(self, tensor, src, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.broadcast(tensor, src, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="broadcast", tensor=tensor, src=src, group=group, async_op=async_op) def all_gather(self, tensor_list, tensor, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.all_gather(tensor_list, tensor, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="all_gather", tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op) - def reduce_scatter_tensor(self, output_tensor,input_tensor, op, group=None, async_op=False): - #todo: ccl version - super(CCLBackend, self).reduce_scatter_tensor(output_tensor,input_tensor, op, group) + def reduce_scatter_tensor(self, output_tensor, input_tensor, op, group=None, async_op=False): + return self.run_collective(name="reduce_scatter_tensor", output_tensor=output_tensor, input_tensor=input_tensor, op=op, group=group) def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False): - #todo: ccl version - super(CCLBackend, self).all_gather_into_tensor(output_tensor, input_tensor, group) + return self.run_collective(name="all_gather_into_tensor", output_tensor=output_tensor, input_tensor=input_tensor, group=group) def all_to_all_single(self, output, input, output_split_sizes, input_split_sizes, group=None, async_op=False): - #todo: ccl version - super(CCLBackend, self).all_to_all_single(output, input, output_split_sizes, input_split_sizes, group) + return self.run_collective(name="all_to_all_single", output=output, input=input, output_split_sizes=output_split_sizes, input_split_sizes=input_split_sizes, group=group) def send(self, tensor, dst, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.send(tensor, dst, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="send", tensor=tensor, dst=dst, group=group, async_op=async_op) def recv(self, tensor, src, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.recv(tensor, src, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="recv", tensor=tensor, src=src, group=group, async_op=async_op) def gather(self, tensor, gather_list, dst, group=None, async_op=False): - #todo: ccl version - super(CCLBackend, self).gather(tensor, gather_list, dst, group) + return self.run_collective(name="gather", tensor=tensor, gather_list=gather_list, dst=dst, group=group) def scatter(self, tensor, gather_list, dst, group=None, async_op=False): - #todo: ccl version - super(CCLBackend, self).scatter(tensor, gather_list, dst, group) + return self.run_collective(name="scatter", tensor=tensor, gather_list=gather_list, dst=dst, group=group) def barrier(self, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.barrier(group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="barrier", group=group, async_op=async_op) def monitored_barrier(self, group=None, timeout=None, wait_all_ranks=False): - #todo: ccl version - super(CCLBackend, self).monitored_barrier(group) + return self.run_collective(name="monitored_barrier", group=group) def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.reduce_scatter(output, input_list, op, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="reduce_scatter", output=output, input_list=input_list, op=op, group=group, async_op=async_op) def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): - group_ranks = self.get_all_ranks_from_group(group) - self.ccl_comm_op.reduce(tensor, dst, op, group_ranks, async_op) - return CCLHandler(self.ccl_comm_op) + return self.run_collective(name="reduce", tensor=tensor, dst=dst, op=op, group=group, async_op=async_op) def new_group(self, ranks): size = len(ranks) From b1d4a7451764c6f9c6a0eed9d2da59ddc484bf62 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Tue, 26 Sep 2023 21:14:11 -0700 Subject: [PATCH 3/5] add coresponding func in cclbackend --- csrc/cpu/comm/ccl.cpp | 38 +++++++++++++++--- deepspeed/comm/ccl.py | 93 +++++++++++++++++++++++++++++-------------- 2 files changed, 95 insertions(+), 36 deletions(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 3291e74fc9c4..65ab7c21065d 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -277,15 +277,15 @@ int world_size = -1; std::set _comm_ids; std::set _colors; -ccl::vector_class _ccl_comms; +std::vector _ccl_comms; +ccl::shared_ptr_class sub_kvs; std::map, int> group_to_comm_id; ccl::communicator& _get_comm_from_group() { return _ccl_comms[0]; } ccl::communicator& _get_comm_from_group(py::object group) { return _ccl_comms[0]; } -ccl::communicator& _get_comm_from_group(std::vector ranks) -{ - if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) - { +ccl::communicator& _get_comm_from_group(std::vector ranks) +{ + if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) { auto id = group_to_comm_id.find(ranks); return _ccl_comms[id->second]; } @@ -412,6 +412,31 @@ py::object new_group(std::vector ranks) << std::endl; } +std::vector get_sub_kvs_addr(bool first) +{ + if (first) { + sub_kvs = ccl::create_main_kvs(); + ccl::kvs::address_type main_addr = sub_kvs->get_address(); + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } else { + ccl::kvs::address_type main_addr; + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } +} + +void initialize_sub_comm(int size, int rank, torch::Tensor& kvs_data, std::vector ranks) +{ + ccl::kvs::address_type main_addr; + if (rank != 0) { + memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); + sub_kvs = ccl::create_kvs(main_addr); + } + _ccl_comms.push_back(ccl::create_communicator(size, rank, sub_kvs)); + group_to_comm_id[ranks] = _ccl_comms.size() - 1; +} + ccl::datatype get_ccl_datatype(c10::ScalarType type) { ccl::datatype ccl_type; @@ -586,7 +611,8 @@ void barrier(std::vector group, bool async_op) std::vector get_available_coll() { - std::vector colls{"broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; + std::vector colls{ + "broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; return colls; } diff --git a/deepspeed/comm/ccl.py b/deepspeed/comm/ccl.py index eef7e6c5a37b..dbc6d481ad30 100644 --- a/deepspeed/comm/ccl.py +++ b/deepspeed/comm/ccl.py @@ -20,14 +20,17 @@ def build_ccl_op(): print(f'DeepSpeed {builder.absolute_name()} built successfully') return ccl_cpp_module + class CCLHandler(): + def __init__(self, ccl_comm_op=None): self.ccl_comm_op = ccl_comm_op - + def wait(self): # backend covered it pass + class CCLBackend(TorchBackend): def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, init_method=None): @@ -50,26 +53,36 @@ def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, i super(CCLBackend, self).broadcast(main_kvs, 0) self.ccl_comm_op.initialize(size, rank, main_kvs) self.initialized = True + self.groups = [tuple(range(self.get_world_size()))] self.available_coll = self.ccl_comm_op.get_available_coll() def is_initialized(self): return self.initialized - + def run_collective(self, name, **kwargs): if name in self.available_coll: kwargs['group'] = self.get_all_ranks_from_group(kwargs['group']) + if 'dst' in kwargs: + kwargs['dst'] = kwargs['group'].index(kwargs['dst']) + if 'src' in kwargs: + kwargs['src'] = kwargs['group'].index(kwargs['src']) func = "self.ccl_comm_op." + name eval(func)(*(kwargs.values())) return CCLHandler(self.ccl_comm_op) else: func = "super(CCLBackend, self)." + name return eval(func)(*(kwargs.values())) - + def all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): use_caching = False if use_caching: match_id = f"{tensor.size()}-{op}" - return self.run_collective(name="all_reduce_caching", tensor=tensor, op=op, match_id=match_id, group=group, async_op=async_op) + return self.run_collective(name="all_reduce_caching", + tensor=tensor, + op=op, + match_id=match_id, + group=group, + async_op=async_op) else: return self.run_collective(name="all_reduce", tensor=tensor, op=op, group=group, async_op=async_op) @@ -77,23 +90,39 @@ def inference_all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=Fal return self.run_collective(name="inference_all_reduce", tensor=tensor, op=op, group=group, async_op=async_op) def broadcast(self, tensor, src, group=None, async_op=False): - return self.run_collective(name="broadcast", tensor=tensor, src=src, group=group, async_op=async_op) + return self.run_collective(name="broadcast", tensor=tensor, src=src, group=group, async_op=async_op) - def all_gather(self, tensor_list, tensor, group=None, async_op=False): - return self.run_collective(name="all_gather", tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op) + def all_gather(self, tensor_list, tensor, group=None, async_op=False): + return self.run_collective(name="all_gather", + tensor_list=tensor_list, + tensor=tensor, + group=group, + async_op=async_op) def reduce_scatter_tensor(self, output_tensor, input_tensor, op, group=None, async_op=False): - return self.run_collective(name="reduce_scatter_tensor", output_tensor=output_tensor, input_tensor=input_tensor, op=op, group=group) + return self.run_collective(name="reduce_scatter_tensor", + output_tensor=output_tensor, + input_tensor=input_tensor, + op=op, + group=group) def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False): - return self.run_collective(name="all_gather_into_tensor", output_tensor=output_tensor, input_tensor=input_tensor, group=group) + return self.run_collective(name="all_gather_into_tensor", + output_tensor=output_tensor, + input_tensor=input_tensor, + group=group) def all_to_all_single(self, output, input, output_split_sizes, input_split_sizes, group=None, async_op=False): - return self.run_collective(name="all_to_all_single", output=output, input=input, output_split_sizes=output_split_sizes, input_split_sizes=input_split_sizes, group=group) + return self.run_collective(name="all_to_all_single", + output=output, + input=input, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group) def send(self, tensor, dst, group=None, async_op=False): return self.run_collective(name="send", tensor=tensor, dst=dst, group=group, async_op=async_op) - + def recv(self, tensor, src, group=None, async_op=False): return self.run_collective(name="recv", tensor=tensor, src=src, group=group, async_op=async_op) @@ -103,43 +132,47 @@ def gather(self, tensor, gather_list, dst, group=None, async_op=False): def scatter(self, tensor, gather_list, dst, group=None, async_op=False): return self.run_collective(name="scatter", tensor=tensor, gather_list=gather_list, dst=dst, group=group) - def barrier(self, group=None, async_op=False): + def barrier(self, group=None, async_op=False): return self.run_collective(name="barrier", group=group, async_op=async_op) - + def monitored_barrier(self, group=None, timeout=None, wait_all_ranks=False): return self.run_collective(name="monitored_barrier", group=group) def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False): - return self.run_collective(name="reduce_scatter", output=output, input_list=input_list, op=op, group=group, async_op=async_op) - - def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): + return self.run_collective(name="reduce_scatter", + output=output, + input_list=input_list, + op=op, + group=group, + async_op=async_op) + + def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): return self.run_collective(name="reduce", tensor=tensor, dst=dst, op=op, group=group, async_op=async_op) - + def new_group(self, ranks): + return super(CCLBackend, self).new_group(ranks) + + def _new_group(self, ranks, group): size = len(ranks) rank = self.get_rank() - if tuple(ranks) in self.groups or rank not in ranks: - return sub_main_kvs = self.ccl_comm_op.get_sub_kvs_addr(rank == ranks[0]) - sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to("xpu:"+str(rank)) - torch_new_group = super(CCLBackend, self).new_group(ranks) - super(CCLBackend, self).broadcast(sub_main_kvs, ranks[0], torch_new_group, False) + sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to("xpu:" + str(rank)) + torch_new_group = group + super(CCLBackend, self).broadcast(sub_main_kvs, ranks[0], group) self.ccl_comm_op.initialize_sub_comm(size, ranks.index(rank), sub_main_kvs, ranks) self.groups.append(tuple(ranks)) - return torch_new_group - + def get_all_ranks_from_group(self, group): if group is None: return list(range(self.get_world_size())) - rank=0 - results=[] + rank = 0 + results = [] try: while True: - results.append(torch.distributed.distributed_c10d._get_global_rank(group, rank)) - rank+=1 + results.append(super(CCLBackend, self).get_global_rank(group, rank)) + rank += 1 except RuntimeError: pass - if tuple(results) not in self.groups: - self.new_group(results) + self._new_group(results, group) return results From 899f55f854d7db9fa5f4807934afdeeb7293d66d Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Wed, 27 Sep 2023 01:49:51 -0700 Subject: [PATCH 4/5] add coresponding pybind func --- csrc/cpu/comm/ccl.cpp | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 65ab7c21065d..0c261b729ba0 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -404,14 +404,6 @@ int next_unique_val(std::set s) } } -py::object new_group(std::vector ranks) -{ - int comm_id = next_unique_val(_comm_ids); - int color = next_unique_val(_colors); - std::cout << "RANK: " << get_rank() << " COMM_ID: " << comm_id << " COLOR: " << color - << std::endl; -} - std::vector get_sub_kvs_addr(bool first) { if (first) { @@ -627,5 +619,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("inference_all_reduce", &inference_all_reduce, "low latency all_reduce implementation"); m.def("all_reduce_caching", &all_reduce_caching, "ccl all_reduce with caching"); m.def("barrier", &barrier, "barrier"); + m.def("initialize_sub_comm", &initialize_sub_comm, "initialize_sub_comm"); + m.def("get_sub_kvs_addr", &get_sub_kvs_addr, "get_sub_kvs_addr"); m.def("get_available_coll", &get_available_coll, "get_available_coll"); } From 1c2836f564012daaefdfc6b373c8a2157bdeade5 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Wed, 11 Oct 2023 18:43:49 -0700 Subject: [PATCH 5/5] change device setting --- deepspeed/comm/ccl.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/deepspeed/comm/ccl.py b/deepspeed/comm/ccl.py index dbc6d481ad30..859e70c7db3b 100644 --- a/deepspeed/comm/ccl.py +++ b/deepspeed/comm/ccl.py @@ -49,7 +49,7 @@ def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, i size = self.get_world_size() rank = self.get_rank() main_kvs = self.ccl_comm_op.get_kvs_addr(rank) - main_kvs = torch.tensor(main_kvs).to(torch.uint8).to(get_accelerator().device_name(rank)) + main_kvs = torch.tensor(main_kvs).to(torch.uint8).to(get_accelerator().current_device_name()) super(CCLBackend, self).broadcast(main_kvs, 0) self.ccl_comm_op.initialize(size, rank, main_kvs) self.initialized = True @@ -156,8 +156,7 @@ def _new_group(self, ranks, group): size = len(ranks) rank = self.get_rank() sub_main_kvs = self.ccl_comm_op.get_sub_kvs_addr(rank == ranks[0]) - sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to("xpu:" + str(rank)) - torch_new_group = group + sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to(get_accelerator().current_device_name()) super(CCLBackend, self).broadcast(sub_main_kvs, ranks[0], group) self.ccl_comm_op.initialize_sub_comm(size, ranks.index(rank), sub_main_kvs, ranks) self.groups.append(tuple(ranks))