已合并
reduce scatter support tensorlist.size != world_size #44440
limuan创建于 25 天前
reduce scatter support tensorlist.size != world_size #44440
已合并
共 5 个文件变更+148-5
| @@ -1,7 +1,7 @@ | |||
| 1 | import os | 1 | import os |
| 2 | import torch | 2 | import torch |
| 3 | import torch.distributed as dist | 3 | import torch.distributed as dist |
| 4 | -import torch_npu | 4 | +import torch_npu # noqa: F401 |
| 5 | 5 | ||
| 6 | 6 | ||
| 7 | def error_size(): | 7 | def error_size(): |
| @@ -12,8 +12,13 @@ def error_size(): | |||
| 12 | dist.init_process_group(backend) | 12 | dist.init_process_group(backend) |
| 13 | rank = dist.get_rank() | 13 | rank = dist.get_rank() |
| 14 | torch.npu.set_device(rank) | 14 | torch.npu.set_device(rank) |
| 15 | - output = torch.tensor(2).npu() | 15 | + ndev = torch.npu.device_count() |
| 16 | - input_list = [torch.tensor(var).npu() for var in range(3)] | 16 | + # output on this rank's device, input_list tensors on a different device |
| 17 | + # -> reduce_scatter rejects input/output residing on different devices. | ||
| 18 | + out_dev = rank | ||
| 19 | + in_dev = (rank + 1) % ndev | ||
| 20 | + output = torch.zeros(4, dtype=torch.float32, device=f"npu:{out_dev}") | ||
| 21 | + input_list = [torch.zeros(4, dtype=torch.float32, device=f"npu:{in_dev}") for _ in range(2)] | ||
| 17 | dist.reduce_scatter(output, input_list) | 22 | dist.reduce_scatter(output, input_list) |
| 18 | 23 | ||
| 19 | 24 | ||
| @@ -117,7 +117,7 @@ class TestMode(TestCase): | |||
| 117 | process.terminate() | 117 | process.terminate() |
| 118 | process.wait() | 118 | process.wait() |
| 119 | self.assertIn( | 119 | self.assertIn( |
| 120 | - "Tensor list input to scatter/gather must match number of collective participants", | 120 | + "Corresponding input/output tensors to reduce_scatter must all reside on the same device", |
| 121 | message | 121 | message |
| 122 | ) | 122 | ) |
| 123 | 123 | ||
| @@ -67,6 +67,51 @@ class HcclReduceScatterTestBase(TestCase): | |||
| 67 | return [input.cpu() for input in inputs] | 67 | return [input.cpu() for input in inputs] |
| 68 | return [input.cpu() * world_size for input in inputs] | 68 | return [input.cpu() * world_size for input in inputs] |
| 69 | 69 | ||
| 70 | + def _numel(self, shape): | ||
| 71 | + n = 1 | ||
| 72 | + for d in shape: | ||
| 73 | + n *= d | ||
| 74 | + return n | ||
| 75 | + | ||
| 76 | + # Expected output for _test_reduce_scatter_lifted. Input convention: the k-th | ||
| 77 | + # flattened element of rank r's i-th tensor is r*10000 + offset_i + k. After SUM, | ||
| 78 | + # global position p (< have) reduces to ws*p + 5000*ws*(ws-1); p >= have is 0. | ||
| 79 | + # rank r takes global [r*out_numel, (r+1)*out_numel). | ||
| 80 | + def _construct_lifted_expected(self, input_shapes, out_shape, world_size): | ||
| 81 | + offsets = [0] | ||
| 82 | + for s in input_shapes: | ||
| 83 | + offsets.append(offsets[-1] + self._numel(s)) | ||
| 84 | + have = offsets[-1] | ||
| 85 | + out_numel = self._numel(out_shape) | ||
| 86 | + base = 5000 * world_size * (world_size - 1) | ||
| 87 | + expected = [] | ||
| 88 | + for r in range(world_size): | ||
| 89 | + vals = [] | ||
| 90 | + for k in range(out_numel): | ||
| 91 | + p = r * out_numel + k | ||
| 92 | + vals.append(world_size * p + base if p < have else 0.0) | ||
| 93 | + expected.append(torch.tensor(vals, dtype=torch.float32).reshape(out_shape)) | ||
| 94 | + return expected | ||
| 95 | + | ||
| 96 | + def _test_multiprocess_lifted(self, fn, init_pg, input_shapes, out_shape, world_size): | ||
| 97 | + ctx = mp.get_context('spawn') | ||
| 98 | + c2p = ctx.Queue(world_size) | ||
| 99 | + p2c = ctx.Queue(world_size) | ||
| 100 | + expected = self._construct_lifted_expected(input_shapes, out_shape, world_size) | ||
| 101 | + ps = [] | ||
| 102 | + for i in range(world_size): | ||
| 103 | + p = ctx.Process(target=fn, args=(i, input_shapes, out_shape, world_size, init_pg, c2p, p2c)) | ||
| 104 | + p.start() | ||
| 105 | + ps.append(p) | ||
| 106 | + for _ in range(world_size): | ||
| 107 | + rank, output = c2p.get() | ||
| 108 | + self.assertEqual(output, expected[rank], | ||
| 109 | + ("rank {} Expect receive tensor {} but got {}.").format(rank, expected[rank], output)) | ||
| 110 | + for _ in range(world_size): | ||
| 111 | + p2c.put(0) | ||
| 112 | + for p in ps: | ||
| 113 | + p.join() | ||
| 114 | + | ||
| 70 | 115 | ||
| 71 | class HcclReduceScatterTest(HcclReduceScatterTestBase): | 116 | class HcclReduceScatterTest(HcclReduceScatterTestBase): |
| 72 | 117 | ||
| @@ -81,6 +126,28 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase): | |||
| 81 | pg.barrier() | 126 | pg.barrier() |
| 82 | p2c.get() | 127 | p2c.get() |
| 83 | 128 | ||
| 129 | + | ||
| 130 | + # pylint:disable=huawei-too-many-arguments | ||
| 131 | + # input_shapes/out_shape are built per-rank with deterministic values so | ||
| 132 | + # the expected result can be computed locally. | ||
| 133 | + def _test_reduce_scatter_lifted(cls, rank, input_shapes, out_shape, world_size, init_pg, c2p, p2c, | ||
| 134 | + reduce_op=dist.ReduceOp.SUM): | ||
| 135 | + pg = init_pg(rank, world_size) | ||
| 136 | + input_list_npu = [] | ||
| 137 | + offset = 0 | ||
| 138 | + for s in input_shapes: | ||
| 139 | + n = 1 | ||
| 140 | + for d in s: | ||
| 141 | + n *= d | ||
| 142 | + vals = torch.arange(offset, offset + n, dtype=torch.float32) + rank * 10000.0 | ||
| 143 | + input_list_npu.append(vals.reshape(s).npu()) | ||
| 144 | + offset += n | ||
| 145 | + output = torch.zeros(out_shape, dtype=torch.float32).npu() | ||
| 146 | + pg.reduce_scatter(output, input_list_npu, reduce_op) | ||
| 147 | + c2p.put((rank, output.cpu())) | ||
| 148 | + pg.barrier() | ||
| 149 | + p2c.get() | ||
| 150 | + | ||
| 84 | 151 | ||
| 85 | # pylint:disable=huawei-too-many-arguments | 152 | # pylint:disable=huawei-too-many-arguments |
| 86 | def _test_reduce_scatter_with_input_internal_format_and_offset(cls, rank, input_list, world_size, init_pg): | 153 | def _test_reduce_scatter_with_input_internal_format_and_offset(cls, rank, input_list, world_size, init_pg): |
| @@ -161,6 +228,44 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase): | |||
| 161 | self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter, | 228 | self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter, |
| 162 | HcclReduceScatterTest._init_dist_hccl, cpu_excepted_result, input_list, world_size) | 229 | HcclReduceScatterTest._init_dist_hccl, cpu_excepted_result, input_list, world_size) |
| 163 | 230 | ||
| 231 | + | ||
| 232 | + | ||
| 233 | + def test_reduce_scatter_single_tensor(self): | ||
| 234 | + # Single-tensor input list (length 1 != world_size). A single long tensor | ||
| 235 | + # of length world_size*out_numel is split evenly across ranks. | ||
| 236 | + ranks = [2] | ||
| 237 | + for world_size in ranks: | ||
| 238 | + out_shape = [4] | ||
| 239 | + input_shapes = [[world_size * 4]] # one tensor, len == world_size*out_numel | ||
| 240 | + self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted, | ||
| 241 | + HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size) | ||
| 242 | + | ||
| 243 | + | ||
| 244 | + | ||
| 245 | + def test_reduce_scatter_input_list_not_equal_world_size(self): | ||
| 246 | + # input_list length differs from world_size: N < ws (zero pad) and N > ws (tail ignore). | ||
| 247 | + ranks = [2] | ||
| 248 | + for world_size in ranks: | ||
| 249 | + out_shape = [4] | ||
| 250 | + # N=1 < ws: have=4 < need=8, trailing rank gets zero pad | ||
| 251 | + self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted, | ||
| 252 | + HcclReduceScatterTest._init_dist_hccl, [[4]], out_shape, world_size) | ||
| 253 | + # N=3 > ws: have=12 > need=8, the 3rd tensor is ignored | ||
| 254 | + self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted, | ||
| 255 | + HcclReduceScatterTest._init_dist_hccl, [[4]] * 3, out_shape, world_size) | ||
| 256 | + | ||
| 257 | + | ||
| 258 | + | ||
| 259 | + def test_reduce_scatter_input_numel_not_equal_output(self): | ||
| 260 | + # N == world_size but per-tensor numel != output numel: have > need (tail ignore). | ||
| 261 | + ranks = [2] | ||
| 262 | + for world_size in ranks: | ||
| 263 | + out_shape = [4] | ||
| 264 | + # each tensor has 6 elements, output 4: have=12 > need=8, tail 2 elements per tensor ignored | ||
| 265 | + input_shapes = [[6]] * world_size | ||
| 266 | + self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted, | ||
| 267 | + HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size) | ||
| 268 | + | ||
| 164 | 269 | ||
| 165 | def test_reduce_scatter_avg(self): | 270 | def test_reduce_scatter_avg(self): |
| 166 | ranks = [2] | 271 | ranks = [2] |
| @@ -240,6 +240,7 @@ class HcclSendRecvDistTest(TestCase): | |||
| 240 | HcclSendRecvDistTest._init_dist_hccl, | 240 | HcclSendRecvDistTest._init_dist_hccl, |
| 241 | world_size) | 241 | world_size) |
| 242 | 242 | ||
| 243 | + | ||
| 243 | 244 | ||
| 244 | def test_send_recv_hccl_dist_with_format(self): | 245 | def test_send_recv_hccl_dist_with_format(self): |
| 245 | self._test_multiprocess( | 246 | self._test_multiprocess( |
| @@ -260,6 +261,7 @@ class HcclSendRecvDistTest(TestCase): | |||
| 260 | torch.randn(2, 2), | 261 | torch.randn(2, 2), |
| 261 | HcclSendRecvDistTest._init_pg_hccl) | 262 | HcclSendRecvDistTest._init_pg_hccl) |
| 262 | 263 | ||
| 264 | + | ||
| 263 | 265 | ||
| 264 | def test_send_recv_hccl_group_with_format(self): | 266 | def test_send_recv_hccl_group_with_format(self): |
| 265 | self._test_multiprocess( | 267 | self._test_multiprocess( |
| @@ -3307,6 +3307,37 @@ std::vector<at::Tensor> flatten_for_scatter_gather( | |||
| 3307 | return flattened; | 3307 | return flattened; |
| 3308 | } | 3308 | } |
| 3309 | 3309 | ||
| 3310 | +// Flatten input tensor lists for reduce_scatter, allowing input list size | ||
| 3311 | +// and per-tensor numel to differ from output (aligning with NCCL). | ||
| 3312 | +std::vector<at::Tensor> flatten_for_reduce_scatter( | ||
| 3313 | + std::vector<std::vector<at::Tensor>>& tensor_lists, | ||
| 3314 | + std::vector<at::Tensor>& outputTensors) | ||
| 3315 | +{ | ||
| 3316 | + if (tensor_lists.size() != outputTensors.size()) { | ||
| 3317 | + TORCH_CHECK( | ||
| 3318 | + false, | ||
| 3319 | + "Tensor list operands to reduce_scatter must have the same length", | ||
| 3320 | + DIST_ERROR(ErrCode::VALUE)); | ||
| 3321 | + } | ||
| 3322 | + const auto num_devices = tensor_lists.size(); | ||
| 3323 | + std::vector<at::Tensor> flattened; | ||
| 3324 | + flattened.resize(num_devices); | ||
| 3325 | + for (const auto i : c10::irange(num_devices)) { | ||
| 3326 | + TORCH_CHECK( | ||
| 3327 | + !tensor_lists[i].empty(), | ||
| 3328 | + "Tensor list operands to reduce_scatter must be non-empty", | ||
| 3329 | + DIST_ERROR(ErrCode::PARAM)); | ||
| 3330 | + if (tensor_lists[i].front().get_device() != outputTensors[i].get_device()) { | ||
 | |||
| 3331 | + TORCH_CHECK( | ||
| 3332 | + false, | ||
| 3333 | + "Corresponding input/output tensors to reduce_scatter must all reside" | ||
| 3334 | + " on the same device", DIST_ERROR(ErrCode::PARAM)); | ||
| 3335 | + } | ||
| 3336 | + flattened[i] = c10d::newLikeFlat(tensor_lists, i); | ||
| 3337 | + } | ||
| 3338 | + return flattened; | ||
| 3339 | +} | ||
| 3340 | + | ||
| 3310 | void nslb_record_end() | 3341 | void nslb_record_end() |
| 3311 | { | 3342 | { |
| 3312 | std::string end_file_path; | 3343 | std::string end_file_path; |
| @@ -5813,7 +5844,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::reduce_scatter( | |||
| 5813 | } | 5844 | } |
| 5814 | bool same_size = check_same_size(inputTensors.back()); | 5845 | bool same_size = check_same_size(inputTensors.back()); |
| 5815 | if (same_size) { | 5846 | if (same_size) { |
| 5816 | - auto inputFlattened = flatten_for_scatter_gather(inputTensors, outputTensors, size_); | 5847 | + auto inputFlattened = flatten_for_reduce_scatter(inputTensors, outputTensors); |
| 5817 | check_npu_tensors_different_devices(inputFlattened); | 5848 | check_npu_tensors_different_devices(inputFlattened); |
| 5818 | std::string functionName = __FUNCTION__; | 5849 | std::string functionName = __FUNCTION__; |
| 5819 | return collective( | 5850 | return collective( |
🟡 Medium Priority
changed line: 第 3324 行
tensor_lists[i].front().get_device()— 新增函数flatten_for_reduce_scatter中，在调用.front()之前没有检查tensor_lists[i]是否为空。受影响行为/契约：原函数
flatten_for_scatter_gather（第 3283 行）通过tensor_lists[i].size() != world_size * num_devices校验，隐式保证了tensor_lists[i]非空（因为world_size >= 1，num_devices >= 1）。新函数flatten_for_reduce_scatter有意去掉了这个 size 校验以支持tensor_lists[i].size() != world_size，但未添加替代的非空检查。失败模式：如果调用方传入空的内层 tensor 列表（即
tensor_lists[i]为空 vector），std::vector::front()行为未定义——通常会导致崩溃或访问非法内存。