已合并
check the size of tensors for alltoall #29731
tonglei创建于 1月19日
check the size of tensors for alltoall #29731
已合并
共 2 个文件变更+6-1
| @@ -79,7 +79,7 @@ class HcomCoalescedManagerTest(TestCase): | |||
| 79 | return expected | 79 | return expected |
| 80 | 80 | ||
| 81 | 81 | ||
| 82 | - @SkipIfNotGteCANNVersion("8.5.0") | 82 | + @SkipIfNotGteCANNVersion("9.0.0") |
| 83 | def test_all_reduce_coalesced_manager_hccl(self): | 83 | def test_all_reduce_coalesced_manager_hccl(self): |
| 84 | ranks = [2] | 84 | ranks = [2] |
| 85 | shape_format = [[np.float32, 2, [2, 3, 16]]] | 85 | shape_format = [[np.float32, 2, [2, 3, 16]]] |
| @@ -5944,6 +5944,11 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::alltoall( | |||
| 5944 | std::vector<at::Tensor>& input_tensors, | 5944 | std::vector<at::Tensor>& input_tensors, |
| 5945 | const c10d::AllToAllOptions& opts) | 5945 | const c10d::AllToAllOptions& opts) |
| 5946 | { | 5946 | { |
| 5947 | + TORCH_CHECK(output_tensors.size() == size_, | ||
| 5948 | + "the size of output_tensors and worldsize must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 5949 | + TORCH_CHECK(output_tensors.size() == input_tensors.size(), | ||
| 5950 | + "the size of input_tensors and output_tensors must equal", DIST_ERROR(ErrCode::PARAM)); | ||
| 5951 | + | ||
| 5947 | auto device = output_tensors[0].device(); | 5952 | auto device = output_tensors[0].device(); |
| 5948 | for (const auto r : c10::irange(output_tensors.size())) { | 5953 | for (const auto r : c10::irange(output_tensors.size())) { |
| 5949 | check_npu_single_tensor(output_tensors[r]); | 5954 | check_npu_single_tensor(output_tensors[r]); |