已合并
[fix] random foreachcopy #4229
yuanjun_li创建于 2月9日
[fix] random foreachcopy #4229
已合并
共 2 个文件变更+61-10
| @@ -25,17 +25,35 @@ | |||
| 25 | 25 | ||
| 26 | namespace op_api { | 26 | namespace op_api { |
| 27 | namespace { | 27 | namespace { |
| 28 | - c10::optional<c10::ScalarType> get_uniform_dtype(const at::TensorList tensors) | 28 | + bool should_group(const at::TensorList tensors, const at::TensorList tensors2) |
| 29 | -{ | 29 | + { |
| 30 | - if (tensors.empty()) return c10::nullopt; | 30 | + if (tensors.empty() && tensors2.empty()) { |
| 31 | - c10::ScalarType first_dtype = tensors[0].scalar_type(); | 31 | + return false; |
| 32 | - for (size_t i = 1; i< tensors.size(); ++i) { | ||
| 33 | - if (tensors[i].scalar_type() != first_dtype) { | ||
| 34 | - return c10::nullopt; | ||
| 35 | } | 32 | } |
| 33 | + | ||
| 34 | + const at::Tensor& base_tensor = !tensors.empty() ? tensors[0] : tensors2[0]; | ||
| 35 | + const auto base_device = base_tensor.device(); | ||
| 36 | + const auto base_dtype = base_tensor.scalar_type(); | ||
| 37 | + | ||
| 38 | + bool has_dtype_mismatch = false; | ||
| 39 | + | ||
| 40 | + auto check_list = [&base_device, &base_dtype, &has_dtype_mismatch](const at::TensorList& list) -> bool { | ||
| 41 | + for (const auto& t : list) { | ||
| 42 | + if (t.device() != base_device) { | ||
| 43 | + return false; | ||
| 44 | + } | ||
| 45 | + if (t.scalar_type() != base_dtype) { | ||
| 46 | + has_dtype_mismatch = true; | ||
| 47 | + } | ||
| 48 | + } | ||
| 49 | + return true; | ||
| 50 | + }; | ||
| 51 | + if (!check_list(tensors) || !check_list(tensors2)) { | ||
| 52 | + return false; | ||
| 53 | + } | ||
| 54 | + | ||
| 55 | + return has_dtype_mismatch; | ||
| 36 | } | 56 | } |
| 37 | - return first_dtype; | ||
| 38 | -} | ||
| 39 | } // namespace anonymous | 57 | } // namespace anonymous |
| 40 | 58 | ||
| 41 | 59 | ||
| @@ -245,7 +263,7 @@ void _foreach_copy_(const at::TensorList self, const at::TensorList src, bool no | |||
| 245 | c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) || | 263 | c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) || |
| 246 | (c10_npu::GetSocVersion() > c10_npu::SocVersion::Ascend910_9391); | 264 | (c10_npu::GetSocVersion() > c10_npu::SocVersion::Ascend910_9391); |
| 247 | 265 | ||
| 248 | - if (get_uniform_dtype(src).has_value()) { | 266 | + if (!should_group(src, self)) { |
| 249 | process_tensor_list_batch(self, src, non_blocking, is_support_nd_out, is_support_batch); | 267 | process_tensor_list_batch(self, src, non_blocking, is_support_nd_out, is_support_batch); |
| 250 | } else { | 268 | } else { |
| 251 | std::unordered_map<c10::ScalarType, std::pair<std::vector<at::Tensor>, std::vector<at::Tensor>>> temp_groups; | 269 | std::unordered_map<c10::ScalarType, std::pair<std::vector<at::Tensor>, std::vector<at::Tensor>>> temp_groups; |
| @@ -0,0 +1,33 @@ | |||
| 1 | +import unittest | ||
| 2 | +import random | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +import numpy as np | ||
| 6 | +import traceback | ||
| 7 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestForeachCopy(TestCase): | ||
| 12 | + | ||
| 13 | + def test_foreach_copy_out_bfloat16_shpae_tensor_num(self): | ||
| 14 | + q1 = torch.rand([2,3,4], device="npu").to(torch.float16) | ||
| 15 | + q2 = torch.rand([2,3,4], device="npu").to(torch.float16) | ||
| 16 | + k1 = torch.zeros([1,2,3], device="cpu", dtype=torch.int64) | ||
| 17 | + k2 = torch.zeros([1,2,3], device="cpu", dtype=torch.int64) | ||
| 18 | + dst_tensors = [] | ||
| 19 | + src_tensors = [] | ||
| 20 | + dst_tensors.append(q1) | ||
| 21 | + dst_tensors.append(k1) | ||
| 22 | + src_tensors.append(q2) | ||
| 23 | + src_tensors.append(k2) | ||
| 24 | + | ||
| 25 | + try: | ||
| 26 | + torch._foreach_copy_(dst_tensors, src_tensors) | ||
| 27 | + except Exception: | ||
| 28 | + traceback.print_exc() | ||
| 29 | + raise AssertionError("foreach copy failed, test won't pass") | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +if __name__ == "__main__": | ||
| 33 | + run_tests() | ||