已合并
feat: add _group_start and _group_end bindings for ProcessGroupHCCL from v2.12.0 #43026
chansinging创建于 7月28日
feat: add _group_start and _group_end bindings for ProcessGroupHCCL from v2.12.0 #43026
已合并
共 2 个文件变更+87-0
| @@ -0,0 +1,83 @@ | |||
| 1 | +import os | ||
| 2 | +import torch | ||
| 3 | +import torch.distributed as dist | ||
| 4 | +from torch.testing._internal.common_utils import run_tests | ||
| 5 | +from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | ||
| 6 | + | ||
| 7 | +import torch_npu # noqa: F401 | ||
| 8 | +from torch_npu.testing.common_distributed import with_comms | ||
| 9 | +from torch.distributed.distributed_c10d import _get_default_group | ||
| 10 | + | ||
| 11 | +os.environ["HCCL_NPU_SOCKET_PORT_RANGE"] = "10000,60000" | ||
| 12 | +os.environ["TIMEOUT"] = "1800" | ||
| 13 | +os.environ["TORCH_TEST_TIMEOUT"] = "1800" | ||
| 14 | + | ||
| 15 | +class TestHcclGroupStartEnd(DTensorTestBase): | ||
| 16 | + | ||
| 17 | + def _get_hccl_backend(self): | ||
| 18 | + if dist.get_backend() != "hccl": | ||
| 19 | + self.skipTest(f"Current backend is {dist.get_backend()}, but this test requires hccl. Skipping.") | ||
| 20 | + pg = _get_default_group() | ||
| 21 | + return pg._get_backend(torch.device(self.device_type)) | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + def test_group_all_reduce_reliability(self): | ||
| 25 | + backend = self._get_hccl_backend() | ||
| 26 | + n = 4 * 1024 * 1024 | ||
| 27 | + | ||
| 28 | + x = torch.ones(n, device=self.device_type) * (self.rank + 1) | ||
| 29 | + y = torch.ones(n, device=self.device_type) * (self.rank + 1) | ||
| 30 | + w = torch.ones(n, device=self.device_type) | ||
| 31 | + | ||
| 32 | + backend._group_start() | ||
| 33 | + work_x = torch.distributed.all_reduce(x, async_op=True) | ||
| 34 | + work_y = torch.distributed.all_reduce(y, async_op=True) | ||
| 35 | + backend._group_end() | ||
| 36 | + | ||
| 37 | + work_x.wait() | ||
| 38 | + self.assertTrue(work_x.is_completed()) | ||
| 39 | + self.assertTrue(work_y.is_completed()) | ||
| 40 | + result_x = torch.matmul(x, w) | ||
| 41 | + | ||
| 42 | + work_y.wait() | ||
| 43 | + result_y = torch.matmul(y, w) | ||
| 44 | + | ||
| 45 | + expected = sum(range(1, self.world_size + 1)) | ||
| 46 | + | ||
| 47 | + self.assertEqual(x, torch.full_like(x, expected)) | ||
| 48 | + self.assertEqual(y, torch.full_like(y, expected)) | ||
| 49 | + self.assertEqual(result_x, expected * n) | ||
| 50 | + self.assertEqual(result_y, expected * n) | ||
| 51 | + | ||
| 52 | + | ||
| 53 | + def test_group_p2p_reliability(self): | ||
| 54 | + if self.world_size < 2: | ||
| 55 | + return | ||
| 56 | + | ||
| 57 | + backend = self._get_hccl_backend() | ||
| 58 | + n = 1024 * 1024 | ||
| 59 | + | ||
| 60 | + x = torch.ones(n, device=self.device_type) * (self.rank + 1) | ||
| 61 | + y = torch.ones(n, device=self.device_type) * (self.rank + 1) | ||
| 62 | + | ||
| 63 | + reqs = [] | ||
| 64 | + | ||
| 65 | + backend._group_start() | ||
| 66 | + if self.rank == 0: | ||
| 67 | + reqs.append(torch.distributed.isend(x, dst=1)) | ||
| 68 | + reqs.append(torch.distributed.isend(y, dst=1)) | ||
| 69 | + elif self.rank == 1: | ||
| 70 | + reqs.append(torch.distributed.irecv(x, src=0)) | ||
| 71 | + reqs.append(torch.distributed.irecv(y, src=0)) | ||
| 72 | + backend._group_end() | ||
| 73 | + | ||
| 74 | + for req in reqs: | ||
| 75 | + req.wait() | ||
| 76 | + self.assertTrue(req.is_completed()) | ||
| 77 | + | ||
| 78 | + if self.rank == 1: | ||
| 79 | + self.assertEqual(x, torch.full_like(x, 1)) | ||
| 80 | + self.assertEqual(y, torch.full_like(y, 1)) | ||
| 81 | + | ||
| 82 | +if __name__ == "__main__": | ||
| 83 | + run_tests() | ||
| @@ -458,6 +458,10 @@ PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs) | |||
| 458 | .def("_add_ephemeral_timeout", | 458 | .def("_add_ephemeral_timeout", |
| 459 | &::c10d_npu::ProcessGroupHCCL::addEphemeralTimeout, | 459 | &::c10d_npu::ProcessGroupHCCL::addEphemeralTimeout, |
| 460 | py::arg("timeout"), | 460 | py::arg("timeout"), |
| 461 | + py::call_guard<py::gil_scoped_release>()) | ||
| 462 | + .def("_group_start", &::c10d_npu::ProcessGroupHCCL::groupStart, | ||
| 463 | + py::call_guard<py::gil_scoped_release>()) | ||
| 464 | + .def("_group_end", &::c10d_npu::ProcessGroupHCCL::groupEnd, | ||
| 461 | py::call_guard<py::gil_scoped_release>()); | 465 | py::call_guard<py::gil_scoped_release>()); |
| 462 | 466 | ||
| 463 | intrusive_ptr_class_<::c10d_npu::ProcessGroupHCCL::Options>( | 467 | intrusive_ptr_class_<::c10d_npu::ProcessGroupHCCL::Options>( |