已合并
check the size of tensors for alltoall #29631
tonglei创建于 1月16日
check the size of tensors for alltoall #29631
已合并
共 2 个文件变更+6-1
| @@ -76,7 +76,7 @@ class HcomCoalescedManagerTest(TestCase): | |||
| 76 | return expected | 76 | return expected |
| 77 | 77 | ||
| 78 | 78 | ||
| 79 | - @SkipIfNotGteCANNVersion("8.5.0") | 79 | + @SkipIfNotGteCANNVersion("9.0.0") |
| 80 | def test_all_reduce_coalesced_manager_hccl(self): | 80 | def test_all_reduce_coalesced_manager_hccl(self): |
| 81 | ranks = [2] | 81 | ranks = [2] |
| 82 | shape_format = [[np.float32, 2, [2, 3, 16]]] | 82 | shape_format = [[np.float32, 2, [2, 3, 16]]] |
| @@ -5989,6 +5989,11 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::alltoall( | |||
| 5989 | std::vector<at::Tensor>& input_tensors, | 5989 | std::vector<at::Tensor>& input_tensors, |
| 5990 | const c10d::AllToAllOptions& opts) | 5990 | const c10d::AllToAllOptions& opts) |
| 5991 | { | 5991 | { |
| 5992 | + TORCH_CHECK(output_tensors.size() == size_, | ||
| 5993 | + "the size of output_tensors and worldsize must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 5994 | + TORCH_CHECK(output_tensors.size() == input_tensors.size(), | ||
| 5995 | + "the size of input_tensors and output_tensors must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 5996 | + | ||
| 5992 | auto device = output_tensors[0].device(); | 5997 | auto device = output_tensors[0].device(); |
| 5993 | for (const auto r : c10::irange(output_tensors.size())) { | 5998 | for (const auto r : c10::irange(output_tensors.size())) { |
| 5994 | check_npu_single_tensor(output_tensors[r]); | 5999 | check_npu_single_tensor(output_tensors[r]); |