support coalescing send/recv-pytorch-AtomGit

Atest/distributed/test_coalesced_manager.py+98-0

@@ -0,0 +1,98 @@
	1	+import unittest
	2	+import os
	3	+import numpy as np
	4	+
	5	+import torch
	6	+import torch.distributed as dist
	7	+from torch.distributed.distributed_c10d import _coalescing_manager
	8	+import torch.multiprocessing as mp
	9	+
	10	+import torch_npu
	11	+from torch_npu.testing.testcase import TestCase, run_tests
	12	+from torch_npu.testing.common_utils import create_common_tensor, SkipIfNotGteCANNVersion
	13	+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
	14	+
	15	+
	16	+class HcomCoalescedManagerTest(TestCase):
	17	+ @classmethod
	18	+ def _init_dist_hccl(cls, rank, world_size):
	19	+ os.environ['MASTER_ADDR'] = '127.0.0.1'
	20	+ os.environ['MASTER_PORT'] = '29500'
	21	+ os.environ['HCCL_WHITELIST_DISABLE'] = '1'
	22	+ torch_npu.npu.set_device(rank)
	23	+ dist.init_process_group(backend='hccl', world_size=world_size, rank=rank)
	24	+ return dist
	25	+
	26	+ @classmethod
	27	+ def _test_all_reduce_coalesced_manager_hccl(cls, rank, input1_list, world_size, init_pg, c2p, reduce_op=dist.ReduceOp.SUM, done_event=None):
	28	+ dist_group = init_pg(rank, world_size)
	29	+ process_group = dist.distributed_c10d._get_default_group()
	30	+ dst = 0
	31	+ device = torch.device(f"npu:{rank:d}")
	32	+ input1_list = [input1.npu() for input1 in input1_list]
	33	+ with _coalescing_manager(group=process_group, device=device, async_ops=True) as cm:
	34	+ for tensor in input1_list:
	35	+ dist.all_reduce(tensor)
	36	+ cm.wait()
	37	+ c2p.put((rank, dst, [input1.cpu() for input1 in input1_list]))
	38	+
	39	+
	40	+ def _test_multiprocess(self, f, init_pg, expected_list, input1_list, world_size, reduce_op=dist.ReduceOp.SUM):
	41	+ ctx = mp.get_context('spawn')
	42	+ c2p = ctx.Queue(world_size)
	43	+ done_event = ctx.Event()
	44	+
	45	+ ps = []
	46	+ for i in range(world_size):
	47	+ p = ctx.Process(
	48	+ target=f,
	49	+ args=(i, [input1.cpu() for input1 in input1_list], world_size, init_pg, c2p, reduce_op, done_event))
	50	+ p.start()
	51	+ ps.append(p)
	52	+
	53	+ for _ in range(world_size):
	54	+ rank, dst, output_list = c2p.get()
	55	+ output_num = len(output_list)
	56	+ if rank == dst:
	57	+ for i in range(output_num):
	58	+ self.assertEqual(output_list[i], expected_list[i],
	59	+ "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
	60	+ rank, world_size, expected_list[i].dtype, expected_list[i].shape, expected_list[i], output_list[i]))
	61	+ done_event.set()
	62	+ for p in ps:
	63	+ p.join()
	64	+
	65	+ def _construct_excepted_result(self, inputs, world_size, dtype=np.float32, reduce_op=dist.ReduceOp.SUM):
	66	+ expected = 0
	67	+ for _ in range(world_size):
	68	+ expected += inputs
	69	+
	70	+ if reduce_op == dist.ReduceOp.AVG:
	71	+ if dtype in [np.int32, np.int8, np.int64, np.uint8]:
	72	+ expected //= world_size
	73	+ else:
	74	+ expected /= world_size
	75	+
	76	+ return expected
	77	+
	78	+ @skipIfUnsupportMultiNPU(2)
	79	+ @SkipIfNotGteCANNVersion("8.5.0")
	80	+ def test_all_reduce_coalesced_manager_hccl(self):
	81	+ ranks = [2]
	82	+ shape_format = [[np.float32, 2, [2, 3, 16]]]
	83	+ op_times = 5
	84	+ input1_list = []
	85	+ expected_list = []
	86	+ for world_size in ranks:
	87	+ for shape in shape_format:
	88	+ for _ in range(op_times):
	89	+ exp_input, input1 = create_common_tensor(shape, -10, 10)
	90	+ expected = self._construct_excepted_result(exp_input, world_size)
	91	+ input1_list.append(input1)
	92	+ expected_list.append(expected)
	93	+ self._test_multiprocess(HcomCoalescedManagerTest._test_all_reduce_coalesced_manager_hccl,
	94	+ HcomCoalescedManagerTest._init_dist_hccl, expected_list, input1_list, world_size)
	95	+
	96	+
	97	+if __name__ == '__main__':
	98	+ run_tests()

Mthird_party/acl/libs/hccl.h+2-0

	@@ -109,4 +109,6 @@ hcclResult_t HcclBatchSendRecv(HcclSendRecvItemDef* sendRecvInfo, u32 itemNum, h
109	hcclResult_t HcclCommInitAll(u32 ndev, s32 devices, hcclComm_t comms);	109	hcclResult_t HcclCommInitAll(u32 ndev, s32 devices, hcclComm_t comms);
110	hcclResult_t HcclCommResume(hcclComm_t comm);	110	hcclResult_t HcclCommResume(hcclComm_t comm);
111	hcclResult_t HcclCommWorkingDevNicSet(HcclComm comm, u32 ranks, bool useBackup, u32 nRanks);	111	hcclResult_t HcclCommWorkingDevNicSet(HcclComm comm, u32 ranks, bool useBackup, u32 nRanks);
		112	+hcclResult_t HcclGroupStart();
		113	+hcclResult_t HcclGroupEnd();
112	}	114	}

Mthird_party/hccl/inc/hccl/hccl.h+10-0

	@@ -504,6 +504,16 @@ extern HcclResult HcclCommDeactivateCommMemory(HcclComm comm, void *virPtr);
504	*/	504	*/
505	extern HcclResult HcclCommWorkingDevNicSet(HcclComm comm, uint32_t ranks, bool useBackup, uint32_t nRanks);	505	extern HcclResult HcclCommWorkingDevNicSet(HcclComm comm, uint32_t ranks, bool useBackup, uint32_t nRanks);
506		506
		507	+/**
		508	+ * @brief Group Start
		509	+ */
		510	+extern HcclResult HcclGroupStart();
		511	+
		512	+/**
		513	+ * @brief Group End
		514	+ */
		515	+extern HcclResult HcclGroupEnd();
		516	+
507	#ifndef OPEN_BUILD_PROJECT	517	#ifndef OPEN_BUILD_PROJECT
508	/**	518	/**
509	* @brief Comm accelerator set/get	519	* @brief Comm accelerator set/get

Mtorch_npu/csrc/distributed/HcclCompile.h+26-0

	@@ -31,6 +31,9 @@ LOAD_FUNCTION(HcclCommRegister)
31	LOAD_FUNCTION(HcclCommDeregister)	31	LOAD_FUNCTION(HcclCommDeregister)
32	LOAD_FUNCTION(HcclCommExchangeMem)	32	LOAD_FUNCTION(HcclCommExchangeMem)
33		33
		34	+REGISTER_LIBRARY(libhcomm)
		35	+REGISTER_FUNCTION(libhcomm, HcclGroupStart)
		36	+REGISTER_FUNCTION(libhcomm, HcclGroupEnd)
34		37
35	extern HcclResult hcclAlltoAllV(const void sendBuf, const void sendCounts, const void *sdispls,	38	extern HcclResult hcclAlltoAllV(const void sendBuf, const void sendCounts, const void *sdispls,
36	HcclDataType sendType, const void recvBuf, const void recvCounts, const void *rdispls,	39	HcclDataType sendType, const void recvBuf, const void recvCounts, const void *rdispls,
	@@ -321,4 +324,27 @@ HcclResult hcclCommExchangeMem(HcclComm comm, void windowHandle, uint32_t peer
321	return ret;	324	return ret;
322	}	325	}
323		326
		327	+HcclResult hcclGroupStart()
		328	+{
		329	+ using hcclGroupStartFunc = HcclResult(*)();
		330	+ static hcclGroupStartFunc func = nullptr;
		331	+ if (func == nullptr) {
		332	+ func = (hcclGroupStartFunc)GET_FUNCTION(libhcomm, HcclGroupStart)
		333	+ }
		334	+ TORCH_CHECK(func, "Failed to find function ", "HcclGroupStart", DIST_ERROR(ErrCode::NOT_FOUND));
		335	+ auto ret = func();
		336	+ return ret;
		337	+}
		338	+
		339	+HcclResult hcclGroupEnd()
		340	+{
		341	+ using hcclGroupEndFunc = HcclResult(*)();
		342	+ static hcclGroupEndFunc func = nullptr;
		343	+ if (func == nullptr) {
		344	+ func = (hcclGroupEndFunc)GET_FUNCTION(libhcomm, HcclGroupEnd)
		345	+ }
		346	+ TORCH_CHECK(func, "Failed to find function ", "HcclGroupEnd", DIST_ERROR(ErrCode::NOT_FOUND));
		347	+ auto ret = func();
		348	+ return ret;
		349	+}
324	} // namespace c10d_npu	350	} // namespace c10d_npu

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.cpp+134-0

	@@ -68,6 +68,7 @@ using hcclUs = std::chrono::steady_clock::time_point;
68		68
69	constexpr int32_t MAX_GROUP_NAME_LEN = 128;	69	constexpr int32_t MAX_GROUP_NAME_LEN = 128;
70	constexpr int32_t NSLB_JOBID_OFFSET = 32;	70	constexpr int32_t NSLB_JOBID_OFFSET = 32;
		71	+static constexpr int CoalActive = 0x01, CoalColl = 0x02, CoalP2P = 0x04;
71		72
72	// HCCL ReduceOp mapping	73	// HCCL ReduceOp mapping
73	std::map<c10d::ReduceOp, HcclReduceOp> hcclOp = {	74	std::map<c10d::ReduceOp, HcclReduceOp> hcclOp = {
	@@ -2720,9 +2721,19 @@ std::vector<std::shared_ptr<HCCLComm>>& ProcessGroupHCCL::createHCCLComm(
2720	std::vector<c10_npu::NPUStream> streamVal;	2721	std::vector<c10_npu::NPUStream> streamVal;
2721	streamVal.reserve(devices.size());	2722	streamVal.reserve(devices.size());
2722		2723
		2724	+ // comms have not been initiated yet, but run HcclGroupEnd before, so end it first
		2725	+ for (const auto i : c10::irange(hcclActiveGroupCounter_)) {
		2726	+ (void)i;
		2727	+ HCCL_CHECK_ERROR(hcclGroupEnd());
		2728	+ }
2723	if (!createHCCLCommEx(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank)) {	2729	if (!createHCCLCommEx(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank)) {
2724	createHCCLCommOrigin(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank);	2730	createHCCLCommOrigin(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank);
2725	}	2731	}
		2732	+ // restart the HcclGroupStart
		2733	+ for (const auto i : c10::irange(hcclActiveGroupCounter_)) {
		2734	+ (void)i;
		2735	+ HCCL_CHECK_ERROR(hcclGroupStart());
		2736	+ }
2726		2737
2727	hcclStreams_.emplace(devicesKey, std::move(streamVal));	2738	hcclStreams_.emplace(devicesKey, std::move(streamVal));
2728		2739
	@@ -3833,6 +3844,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
3833		3844
3834	const auto devices = getDevice(inputs);	3845	const auto devices = getDevice(inputs);
3835	auto key = getKeyFromDevice(devices);	3846	auto key = getKeyFromDevice(devices);
		3847	+ NPU_CHECK_ERROR(c10_npu::SetDevice(devices[0].index()));
3836	HcclCommConfig config = createHcclCommConfigWithOptions();	3848	HcclCommConfig config = createHcclCommConfigWithOptions();
3837	std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);	3849	std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);
3838		3850
	@@ -3853,6 +3865,29 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
3853	// For sync operations, we skip syncStreams since we're using the current stream directly	3865	// For sync operations, we skip syncStreams since we're using the current stream directly
3854	}	3866	}
3855		3867
		3868	+ // First let HCCL streams wait for input tensors allocation streams
		3869	+ if (coalescing_state_ & CoalActive) {
		3870	+ coalescing_state_ \|= CoalColl;
		3871	+ if (coalescedDevice_.index() < 0) {
		3872	+ coalescedDevice_ = devices[0];
		3873	+ } else {
		3874	+ for (const auto& device : devices) {
		3875	+ TORCH_CHECK(
		3876	+ coalescedDevice_.index() == device.index(),
		3877	+ "Expecting same device across coalesced P2P operations. "
		3878	+ "Got device ", device.index(), " but expected ", coalescedDevice_.index());
		3879	+ }
		3880	+ }
		3881	+ if (coalescedComm_ == nullptr) {
		3882	+ coalescedComm_ = hcclComms[0];
		3883	+ } else {
		3884	+ // For multi-device, we check if the first comm matches
		3885	+ TORCH_CHECK(
		3886	+ coalescedComm_ == hcclComms[0],
		3887	+ "Expecting same communicator across coalesced P2P operations.");
		3888	+ }
		3889	+ }
		3890	+
3856	// Work itself will create the events on all NPUs of tensors	3891	// Work itself will create the events on all NPUs of tensors
3857	auto work = initWork(devices, rank_, opType);	3892	auto work = initWork(devices, rank_, opType);
3858	// Store references to outputs to be used by WorkHCCL::result and operator<<.	3893	// Store references to outputs to be used by WorkHCCL::result and operator<<.
	@@ -5486,6 +5521,105 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::barrier(const c10d::BarrierOpti
5486	return work;	5521	return work;
5487	}	5522	}
5488		5523
		5524	+void ProcessGroupHCCL::startCoalescing()
		5525	+{
		5526	+ coalescedDevice_.set_index(-1);
		5527	+ coalescedComm_ = nullptr;
		5528	+ coalescedTensors_.clear();
		5529	+ coalescing_state_ \|= CoalActive;
		5530	+ groupStart();
		5531	+}
		5532	+
		5533	+// `optype` is for specifying a composite optype, such as ALLGATHER and
		5534	+// REDUCE_SCATTER
		5535	+c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::endCoalescing(c10d::OpType optype)
		5536	+{
		5537	+ if (coalescedComm_ == nullptr) {
		5538	+ // There is no actual work being coalesced, return here
		5539	+ groupEnd();
		5540	+ coalescing_state_ = 0;
		5541	+ return nullptr;
		5542	+ }
		5543	+ TORCH_CHECK(
		5544	+ coalescedDevice_.index() >= 0,
		5545	+ "Something went wrong. Did you call end_coalescing before start_coalescing?");
		5546	+
		5547	+ // `coalescedComm_` should have same set of comms across collectives
		5548	+ auto comm = coalescedComm_;
		5549	+ // `coalescedDevice_` should have same set of devices across collectives
		5550	+ auto device = coalescedDevice_;
		5551	+ std::vector<at::Device> devices = {device};
		5552	+
		5553	+ // `getKeyFromDevice` is how we get keys for both collectives and batch P2P
		5554	+ const auto key = getKeyFromDevice(devices);
		5555	+ auto& hcclStreams = hcclStreams_[key];
		5556	+ c10_npu::NPUStream& hcclStream = hcclStreams[0];
		5557	+ auto opProfilerTitle = optype != c10d::OpType::COALESCED
		5558	+ ? "hccl:" + opTypeToString(optype) + "_coalesced"
		5559	+ : "hccl:coalesced";
		5560	+
		5561	+ // Create Work object
		5562	+ c10_npu::CaptureStatus capture_status = c10_npu::currentStreamCaptureStatusMayInitCtx();
		5563	+ bool enqueue = (coalescing_state_) && capture_status == c10_npu::CaptureStatus::None;
		5564	+ auto work = initWork(
		5565	+ std::vector<c10::Device>{device},
		5566	+ rank_,
		5567	+ optype,
		5568	+ opProfilerTitle.c_str(),
		5569	+ {},
		5570	+ {},
		5571	+ enqueue);
		5572	+ work->hcclComms_[0] = comm;
		5573	+ work->blockingWait_ = blockingWait_;
		5574	+ work->opTimeout_ = options_->timeout;
		5575	+
		5576	+ // Record start before hcclGroupEnd
		5577	+ if (desyncDebug_) {
		5578	+ (*(work->hcclStartEvents_))[0].record(hcclStream);
		5579	+ }
		5580	+ // Set device before hcclGroupEnd
		5581	+ NPU_CHECK_ERROR(c10_npu::SetDevice(device.index()));
		5582	+ groupEnd();
		5583	+
		5584	+ if (enqueue) {
		5585	+ c10_npu::NPUGraph::inc_pending_event_queries();
		5586	+ workEnqueue(work);
		5587	+ }
		5588	+ {
		5589	+ c10_npu::NPUMultiStreamGuard guard(hcclStreams);
		5590	+ work->future_ = c10::make_intrusive<at::ivalue::Future>(
		5591	+ c10::ListType::create(c10::TensorType::get()),
		5592	+ devices);
		5593	+ work->future_->markCompleted(at::IValue(std::vector<at::Tensor>{}));
		5594	+ }
		5595	+
		5596	+ // Reset coalescing state
		5597	+ coalescing_state_ = 0;
		5598	+ coalescedComm_ = nullptr;
		5599	+ coalescedTensors_.clear();
		5600	+ // If in async mode, return work; otherwise, kernel is enqueued on current
		5601	+ // stream, no need to return work
		5602	+ return work;
		5603	+}
		5604	+
		5605	+void ProcessGroupHCCL::groupStart()
		5606	+{
		5607	+ HCCL_CHECK_ERROR(hcclGroupStart());
		5608	+ ++hcclActiveGroupCounter_;
		5609	+}
		5610	+
		5611	+void ProcessGroupHCCL::groupEnd()
		5612	+{
		5613	+ HCCL_CHECK_ERROR(hcclGroupEnd());
		5614	+ --hcclActiveGroupCounter_;
		5615	+}
		5616	+
		5617	+c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::endCoalescing()
		5618	+{
		5619	+ // Default OpType to COALESCED if not specified
		5620	+ return endCoalescing(c10d::OpType::COALESCED);
		5621	+}
		5622	+
5489	c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::gather(	5623	c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::gather(
5490	std::vector<std::vector<at::Tensor>>& /* unused */,	5624	std::vector<std::vector<at::Tensor>>& /* unused */,
5491	std::vector<at::Tensor>& /* unused */,	5625	std::vector<at::Tensor>& /* unused */,

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.hpp+25-0

	@@ -552,6 +552,19 @@ public:
552	{	552	{
553	return std::string(HCCL_BACKEND_NAME);	553	return std::string(HCCL_BACKEND_NAME);
554	}	554	}
		555	+
		556	+ bool supportsCoalescing() const override
		557	+ {
		558	+ return true;
		559	+ }
		560	+
		561	+ void startCoalescing() override;
		562	+
		563	+ c10::intrusive_ptr<c10d::Work> endCoalescing() override;
		564	+
		565	+ // For specifying a composite optype, such as ALLGATHER and REDUCE_SCATTER
		566	+ c10::intrusive_ptr<c10d::Work> endCoalescing(c10d::OpType optype);
		567	+
555	c10::intrusive_ptr<c10d::Work> broadcast(	568	c10::intrusive_ptr<c10d::Work> broadcast(
556	std::vector<at::Tensor>& tensors,	569	std::vector<at::Tensor>& tensors,
557	const c10d::BroadcastOptions& opts = c10d::BroadcastOptions()) override;	570	const c10d::BroadcastOptions& opts = c10d::BroadcastOptions()) override;
	@@ -652,6 +665,10 @@ public:
652	int srcRank,	665	int srcRank,
653	int tag) override;	666	int tag) override;
654		667
		668	+ void groupStart();
		669	+
		670	+ void groupEnd();
		671	+
655	c10::intrusive_ptr<c10d::Work> recvAnysource(	672	c10::intrusive_ptr<c10d::Work> recvAnysource(
656	std::vector<at::Tensor>& tensors,	673	std::vector<at::Tensor>& tensors,
657	int tag) override;	674	int tag) override;
	@@ -965,6 +982,14 @@ protected:
965	// Device Indexes used for all collectives in this group	982	// Device Indexes used for all collectives in this group
966	std::set<int> usedDeviceIdxs_;	983	std::set<int> usedDeviceIdxs_;
967		984
		985	+ int coalescing_state_ = 0;
		986	+
		987	+ at::Device coalescedDevice_ = at::Device("npu");
		988	+
		989	+ std::shared_ptr<HCCLComm> coalescedComm_ = nullptr;
		990	+
		991	+ TensorShelf coalescedTensors_;
		992	+
968	// map from the key: "group name + pg counter (ID)" to the	993	// map from the key: "group name + pg counter (ID)" to the
969	// HCCL Master ID count. This needs to be group and pg specific	994	// HCCL Master ID count. This needs to be group and pg specific
970		995