已合并
check the size of tensors for alltoall #29942
tonglei创建于 1月22日
check the size of tensors for alltoall #29942
已合并
共 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]]] |
| @@ -6005,6 +6005,11 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::alltoall( | |||
| 6005 | std::vector<at::Tensor>& input_tensors, | 6005 | std::vector<at::Tensor>& input_tensors, |
| 6006 | const c10d::AllToAllOptions& opts) | 6006 | const c10d::AllToAllOptions& opts) |
| 6007 | { | 6007 | { |
| 6008 | + TORCH_CHECK(output_tensors.size() == size_, | ||
| 6009 | + "the size of output_tensors and worldsize must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 6010 | + TORCH_CHECK(output_tensors.size() == input_tensors.size(), | ||
| 6011 | + "the size of input_tensors and output_tensors must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 6012 | + | ||
| 6008 | auto device = output_tensors[0].device(); | 6013 | auto device = output_tensors[0].device(); |
| 6009 | for (const auto r : c10::irange(output_tensors.size())) { | 6014 | for (const auto r : c10::irange(output_tensors.size())) { |
| 6010 | check_npu_single_tensor(output_tensors[r]); | 6015 | check_npu_single_tensor(output_tensors[r]); |