Add ProcessGroupHCCL::getCollNpuStreamId() to get or create NPU stream by device-pytorch-AtomGit

Atest/distributed/test_get_coll_npu_stream_id.py+133-0

@@ -0,0 +1,133 @@
	1	+import os
	2	+from unittest.mock import patch
	3	+
	4	+import torch
	5	+import torch.distributed as dist
	6	+import torch.multiprocessing as mp
	7	+from torch.distributed.distributed_c10d import _world
	8	+
	9	+import torch_npu
	10	+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
	11	+from torch_npu.testing.testcase import TestCase, run_tests
	12	+
	13	+
	14	+class CollNpuStreamIdTest(TestCase):
	15	+ @classmethod
	16	+ def _init_dist_hccl(cls, rank, world_size):
	17	+ os.environ['MASTER_ADDR'] = '127.0.0.1'
	18	+ os.environ['MASTER_PORT'] = '29502'
	19	+ os.environ['HCCL_WHITELIST_DISABLE'] = '1'
	20	+ torch_npu.npu.set_device(rank)
	21	+ dist.init_process_group(backend='hccl', world_size=world_size, rank=rank)
	22	+ return dist
	23	+
	24	+ @classmethod
	25	+ def _test_coll_npu_stream_id_after_collective(cls, rank, world_size, init_pg, c2p, p2c):
	26	+ dist_group = init_pg(rank, world_size)
	27	+
	28	+ backend = _world.default_pg._get_backend(torch.device('npu'))
	29	+ device = torch.device(f'npu:{rank}')
	30	+
	31	+ input1 = torch.tensor([1]).npu()
	32	+ dist_group.all_reduce(input1)
	33	+
	34	+ stream_id = backend.get_coll_stream_id(device)
	35	+
	36	+ assert0 = (stream_id != -1)
	37	+ if stream_id != -1:
	38	+ stream = torch.npu.Stream(stream_id=stream_id, device_type=20, device_index=device.index)
	39	+ assert1 = (stream.npu_stream is not None)
	40	+ else:
	41	+ assert1 = False
	42	+
	43	+ c2p.put(assert0 and assert1)
	44	+ p2c.get()
	45	+
	46	+ @classmethod
	47	+ def _test_coll_npu_stream_id_idempotent(cls, rank, world_size, init_pg, c2p, p2c):
	48	+ dist_group = init_pg(rank, world_size)
	49	+
	50	+ backend = _world.default_pg._get_backend(torch.device('npu'))
	51	+ device = torch.device(f'npu:{rank}')
	52	+
	53	+ input1 = torch.tensor([1]).npu()
	54	+ dist_group.all_reduce(input1)
	55	+
	56	+ stream_id_first = backend.get_coll_stream_id(device)
	57	+ stream_id_second = backend.get_coll_stream_id(device)
	58	+
	59	+ assert0 = (stream_id_first == stream_id_second)
	60	+
	61	+ c2p.put(assert0)
	62	+ p2c.get()
	63	+
	64	+ @classmethod
	65	+ def _test_coll_npu_stream_id_before_collective(cls, rank, world_size, init_pg, c2p, p2c):
	66	+ dist_group = init_pg(rank, world_size)
	67	+
	68	+ backend = _world.default_pg._get_backend(torch.device('npu'))
	69	+ device = torch.device(f'npu:{rank}')
	70	+
	71	+ stream_id = backend.get_coll_stream_id(device)
	72	+
	73	+ assert0 = (stream_id != -1)
	74	+ if stream_id != -1:
	75	+ stream = torch.npu.Stream(stream_id=stream_id, device_type=20, device_index=device.index)
	76	+ assert1 = (stream.npu_stream is not None)
	77	+ else:
	78	+ assert1 = False
	79	+
	80	+ c2p.put(assert0 and assert1)
	81	+ p2c.get()
	82	+
	83	+ def _run_multiprocess_test(self, test_func, init_pg, world_size):
	84	+ ctx = mp.get_context('spawn')
	85	+ c2p = ctx.Queue(world_size)
	86	+ p2c = ctx.Queue(world_size)
	87	+
	88	+ ps = []
	89	+ for rank in range(world_size):
	90	+ p = ctx.Process(target=test_func, args=(rank, world_size, init_pg, c2p, p2c))
	91	+ p.start()
	92	+ ps.append(p)
	93	+
	94	+ for _ in range(world_size):
	95	+ output = c2p.get()
	96	+ self.assertEqual(True, output)
	97	+
	98	+ for _ in range(world_size):
	99	+ p2c.put(0)
	100	+
	101	+ for p in ps:
	102	+ p.join()
	103	+
	104	+ @skipIfUnsupportMultiNPU(2)
	105	+ def test_coll_npu_stream_id_after_collective(self):
	106	+ with patch.dict(os.environ, {"P2P_HCCL_BUFFSIZE": "0"}):
	107	+ self._run_multiprocess_test(
	108	+ CollNpuStreamIdTest._test_coll_npu_stream_id_after_collective,
	109	+ CollNpuStreamIdTest._init_dist_hccl,
	110	+ 2
	111	+ )
	112	+
	113	+ @skipIfUnsupportMultiNPU(2)
	114	+ def test_coll_npu_stream_id_idempotent(self):
	115	+ with patch.dict(os.environ, {"P2P_HCCL_BUFFSIZE": "0"}):
	116	+ self._run_multiprocess_test(
	117	+ CollNpuStreamIdTest._test_coll_npu_stream_id_idempotent,
	118	+ CollNpuStreamIdTest._init_dist_hccl,
	119	+ 2
	120	+ )
	121	+
	122	+ @skipIfUnsupportMultiNPU(2)
	123	+ def test_coll_npu_stream_id_before_collective(self):
	124	+ with patch.dict(os.environ, {"P2P_HCCL_BUFFSIZE": "0"}):
	125	+ self._run_multiprocess_test(
	126	+ CollNpuStreamIdTest._test_coll_npu_stream_id_before_collective,
	127	+ CollNpuStreamIdTest._init_dist_hccl,
	128	+ 2
	129	+ )
	130	+
	131	+
	132	+if __name__ == '__main__':
	133	+ run_tests()

Mtorch_npu/csrc/distributed/Init.cpp+2-0

	@@ -418,6 +418,8 @@ PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs)
418	.def("_get_stream_id", &::c10d_npu::ProcessGroupHCCL::getStreamId,	418	.def("_get_stream_id", &::c10d_npu::ProcessGroupHCCL::getStreamId,
419	py::arg("p2p") = false,	419	py::arg("p2p") = false,
420	py::arg("peer") = -1)	420	py::arg("peer") = -1)
		421	+ .def("get_coll_stream_id", &::c10d_npu::ProcessGroupHCCL::getCollNpuStreamId,
		422	+ py::arg("device"))
421	.def("get_p2p_stream_id", &::c10d_npu::ProcessGroupHCCL::getP2PStreamId,	423	.def("get_p2p_stream_id", &::c10d_npu::ProcessGroupHCCL::getP2PStreamId,
422	py::arg("device"),	424	py::arg("device"),
423	py::arg("peer"),	425	py::arg("peer"),

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

	@@ -2930,6 +2930,20 @@ int64_t ProcessGroupHCCL::getStreamId(bool p2p, int peer)
2930	return hcclStreams_[key][0].id();	2930	return hcclStreams_[key][0].id();
2931	}	2931	}
2932		2932
		2933	+int64_t ProcessGroupHCCL::getCollNpuStreamId(at::Device device)
		2934	+{
		2935	+ const auto key = getKeyFromDevice({device});
		2936	+ if (hcclStreams_.find(key) == hcclStreams_.end() \|\| hcclStreams_[key].empty()) {
		2937	+ bool force_high = c10d::getCvarBool(TORCH_HCCL_HIGH_PRIORITY, false);
		2938	+ auto streamVal = c10_npu::getStreamFromPool(
		2939	+ options_->is_high_priority_stream \|\| force_high, device.index());
		2940	+ hcclStreams_.emplace(key, std::vector<c10_npu::NPUStream>{streamVal});
		2941	+ return streamVal.id();
		2942	+ }
		2943	+ auto hcclStream = hcclStreams_[key][0];
		2944	+ return hcclStream.id();
		2945	+}
		2946	+
2933	int64_t ProcessGroupHCCL::getP2PStreamId(	2947	int64_t ProcessGroupHCCL::getP2PStreamId(
2934	at::Device device,	2948	at::Device device,
2935	int peer,	2949	int peer,

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.hpp+4-2

	@@ -112,7 +112,7 @@ static std::vector<std::string> TORCH_HCCL_HIGH_PRIORITY = {
112	// A struct to hold the latest status of the process group.	112	// A struct to hold the latest status of the process group.
113	struct ProcessGroupStatus {	113	struct ProcessGroupStatus {
114	// the sequential number of the last collective enqueued into workMetaList_	114	// the sequential number of the last collective enqueued into workMetaList_
115	- // This is useful for indentifying a rank that has not join a collective	115	+ // This is useful for identifying a rank that has not join a collective
116	// initialized to be -1 to indicate no collective has been enqueued	116	// initialized to be -1 to indicate no collective has been enqueued
117	int64_t lastEnqueuedSeq{-1};	117	int64_t lastEnqueuedSeq{-1};
118	// the sequential number of the last collective started as the kernel	118	// the sequential number of the last collective started as the kernel
	@@ -796,6 +796,8 @@ public:
796		796
797	int64_t getStreamId(bool p2p, int peer);	797	int64_t getStreamId(bool p2p, int peer);
798		798
		799	+ int64_t getCollNpuStreamId(at::Device device);
		800	+
799	int64_t getP2PStreamId(at::Device device, int peer, int is_batched);	801	int64_t getP2PStreamId(at::Device device, int peer, int is_batched);
800		802
801	void windowRegisterAndExchange(int64_t windowSize, std::vector<uint32_t>& peerRanks);	803	void windowRegisterAndExchange(int64_t windowSize, std::vector<uint32_t>& peerRanks);
	@@ -1018,7 +1020,7 @@ protected:
1018	std::atomic<bool> collectiveDebugInfoMode_;	1020	std::atomic<bool> collectiveDebugInfoMode_;
1019		1021
1020	// This is the signal from watchdog threads to indicate whether the monitor	1022	// This is the signal from watchdog threads to indicate whether the monitor
1021	- // thread should dump. Making it static so that it is accessiable from all the	1023	+ // thread should dump. Making it static so that it is accessible from all the
1022	// PGs. With this flag, monitor thread would dump debug info under any one of	1024	// PGs. With this flag, monitor thread would dump debug info under any one of
1023	// the 3 conditions: 1: this flag is set to true by the watchdog thread when	1025	// the 3 conditions: 1: this flag is set to true by the watchdog thread when
1024	// it detects a timeout. 2: timeout signal is received from	1026	// it detects a timeout. 2: timeout signal is received from