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

已合并

Add ProcessGroupHCCL::getCollNpuStreamId() to get or create NPU stream by device #34700

limuan创建于 4月29日从limuan/pytorch： collstreamid/v2.9.0 合入到 Ascend/pytorch： v2.9.0

Add ProcessGroupHCCL::getCollNpuStreamId() to get or create NPU stream by device #34700

已合并

limuan创建于 4月29日

从limuan/pytorch collstreamid/v2.9.0合入到Ascend/pytorch v2.9.0

共 4 个文件变更+151-0

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

	@@ -2943,6 +2943,20 @@ int64_t ProcessGroupHCCL::getStreamId(bool p2p, int peer)
2943	return hcclStreams_[key][0].id();	2943	return hcclStreams_[key][0].id();
2944	}	2944	}
2945		2945
		2946	+int64_t ProcessGroupHCCL::getCollNpuStreamId(at::Device device)
		2947	+{
		2948	+ const auto key = getKeyFromDevice({device});
		2949	+ if (hcclStreams_.find(key) == hcclStreams_.end() \|\| hcclStreams_[key].empty()) {
		2950	+ bool force_high = c10d::getCvarBool(TORCH_HCCL_HIGH_PRIORITY, false);
		2951	+ auto streamVal = c10_npu::getStreamFromPool(
		2952	+ options_->is_high_priority_stream \|\| force_high, device.index());
		2953	+ hcclStreams_.emplace(key, std::vector<c10_npu::NPUStream>{streamVal});
		2954	+ return streamVal.id();
		2955	+ }
		2956	+ auto hcclStream = hcclStreams_[key][0];
		2957	+ return hcclStream.id();
		2958	+}
		2959	+
2946	int64_t ProcessGroupHCCL::getP2PStreamId(	2960	int64_t ProcessGroupHCCL::getP2PStreamId(
2947	at::Device device,	2961	at::Device device,
2948	int peer,	2962	int peer,

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

	@@ -803,6 +803,8 @@ public:
803		803
804	int64_t getStreamId(bool p2p, int peer);	804	int64_t getStreamId(bool p2p, int peer);
805		805
		806	+ int64_t getCollNpuStreamId(at::Device device);
		807	+
806	int64_t getP2PStreamId(at::Device device, int peer, int is_batched);	808	int64_t getP2PStreamId(at::Device device, int peer, int is_batched);
807		809
808	void windowRegisterAndExchange(int64_t windowSize, std::vector<uint32_t>& peerRanks);	810	void windowRegisterAndExchange(int64_t windowSize, std::vector<uint32_t>& peerRanks);