已合并
add multi_devices_single_process UT #11954
tangpeiqi96创建于 2024年5月23日
add multi_devices_single_process UT #11954
已合并
从refs/pull/11954/head合入到master
共 1 个文件变更+103-0
| @@ -1,9 +1,11 @@ | |||
| 1 | import os | 1 | import os |
| 2 | import torch | 2 | import torch |
| 3 | +import numpy as np | ||
| 3 | 4 | ||
| 4 | import torch_npu | 5 | import torch_npu |
| 5 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | 6 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU |
| 6 | from torch_npu.testing.testcase import TestCase, run_tests | 7 | from torch_npu.testing.testcase import TestCase, run_tests |
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 7 | 9 | ||
| 8 | os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False' | 10 | os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False' |
| 9 | 11 | ||
| @@ -89,5 +91,106 @@ class TestNpu(TestCase): | |||
| 89 | self._test_module() | 91 | self._test_module() |
| 90 | 92 | ||
| 91 | 93 | ||
| 94 | +class TestOp(TestCase): | ||
| 95 | + | ||
| 96 | + def _cpu_op_exec(self, input1): | ||
| 97 | + output = torch.abs(input1) | ||
| 98 | + output = output.cpu().numpy() | ||
| 99 | + return output | ||
| 100 | + | ||
| 101 | + def _npu_op_exec(self, input1): | ||
| 102 | + output = torch.abs(input1) | ||
| 103 | + output = output.cpu().numpy() | ||
| 104 | + return output | ||
| 105 | + | ||
| 106 | + def _test_abs(self, device="npu:1"): | ||
| 107 | + torch.npu.set_device(0) | ||
| 108 | + cpu_input = torch.Tensor([1, -2, -10]) | ||
| 109 | + npu_input = cpu_input.to(device) | ||
| 110 | + cpu_output = self._cpu_op_exec(cpu_input) | ||
| 111 | + npu_output = self._npu_op_exec(npu_input) | ||
| 112 | + self.assertRtolEqual(cpu_output, npu_output) | ||
| 113 | + | ||
| 114 | + def _test_isfinite(self, device="npu:1"): | ||
| 115 | + torch.npu.set_device(0) | ||
| 116 | + x = torch.Tensor([1, 2, -10]).to(device) | ||
| 117 | + output = torch.isfinite(x) | ||
| 118 | + self.assertTrue(output.all()) | ||
| 119 | + | ||
| 120 | + def _test_unique_dim(self, device="npu:1", dtype=torch.float): | ||
| 121 | + torch.npu.set_device(0) | ||
| 122 | + self.assertFalse(hasattr(torch, "unique_dim")) | ||
| 123 | + | ||
| 124 | + x = torch.tensor([[[1., 1.], | ||
| 125 | + [0., 1.], | ||
| 126 | + [2., 1.], | ||
| 127 | + [0., 1.]], | ||
| 128 | + [[1., 1.], | ||
| 129 | + [0., 1.], | ||
| 130 | + [2., 1.], | ||
| 131 | + [0., 1.]]], | ||
| 132 | + dtype=dtype, | ||
| 133 | + device=device) | ||
| 134 | + expected_unique_dim0 = torch.tensor([[[1., 1.], | ||
| 135 | + [0., 1.], | ||
| 136 | + [2., 1.], | ||
| 137 | + [0., 1.]]], | ||
| 138 | + dtype=dtype, | ||
| 139 | + device=device) | ||
| 140 | + expected_inverse_dim0 = torch.tensor([0, 0]) | ||
| 141 | + expected_counts_dim0 = torch.tensor([2]) | ||
| 142 | + | ||
| 143 | + x_unique, x_inverse, x_counts = torch.unique( | ||
| 144 | + x, | ||
| 145 | + return_inverse=True, | ||
| 146 | + return_counts=True, | ||
| 147 | + dim=0) | ||
| 148 | + self.assertEqual(expected_unique_dim0, x_unique) | ||
| 149 | + self.assertEqual(expected_inverse_dim0, x_inverse) | ||
| 150 | + self.assertEqual(expected_counts_dim0, x_counts) | ||
| 151 | + | ||
| 152 | + def _supported_op_exec(self, query_states1, past_key, past_value, head_dim): | ||
| 153 | + attn_weights1 = torch.matmul(query_states1, past_key.transpose(2, 3)) / 0.0078125 | ||
| 154 | + attn_weights1 = torch.max(attn_weights1, torch.full( | ||
| 155 | + (1, 1), torch.finfo(attn_weights1.dtype).min, device=attn_weights1.device)) | ||
| 156 | + attn_weights1 = torch.nn.functional.softmax(attn_weights1, dim=-1, dtype=torch.float32).to(query_states1.dtype) | ||
| 157 | + attn_output1 = torch.matmul(attn_weights1, past_value) | ||
| 158 | + return attn_output1 | ||
| 159 | + | ||
| 160 | + def _custom_op_exec(self, query, key, value, head_dim): | ||
| 161 | + scale = 1 / 0.0078125 | ||
| 162 | + return torch_npu.npu_prompt_flash_attention( | ||
| 163 | + query, key, value, num_heads=32, input_layout="BNSD", scale_value=scale, pre_tokens=65535, next_tokens=65535, sparse_mode=0) | ||
| 164 | + | ||
| 165 | + | ||
| 166 | + def _test_npu_prompt_flash_attention(self, device="npu:1"): | ||
| 167 | + torch.npu.set_device(0) | ||
| 168 | + query = torch.randn(1, 32, 2048, 128, dtype=torch.float16).to(device) | ||
| 169 | + key = torch.randn(1, 32, 2048, 128, dtype=torch.float16).to(device) | ||
| 170 | + value = torch.randn(1, 32, 2048, 128, dtype=torch.float16).to(device) | ||
| 171 | + | ||
| 172 | + head_dim = 128 | ||
| 173 | + | ||
| 174 | + supported_output = self._supported_op_exec(query, key, value, head_dim) | ||
| 175 | + custom_output = self._custom_op_exec(query, key, value, head_dim) | ||
| 176 | + self.assertRtolEqual(supported_output, custom_output) | ||
| 177 | + | ||
| 178 | + | ||
| 179 | + def test_aclop_op_with_multi_device(self): | ||
| 180 | + torch.npu.set_compile_mode(jit_compile=True) | ||
| 181 | + self._test_abs() | ||
| 182 | + self._test_isfinite() | ||
| 183 | + self._test_unique_dim() | ||
| 184 | + self._test_npu_prompt_flash_attention() | ||
| 185 | + | ||
| 186 | + | ||
| 187 | + def test_opapi_op_with_multi_device(self): | ||
| 188 | + torch.npu.set_compile_mode(jit_compile=False) | ||
| 189 | + self._test_abs() | ||
| 190 | + self._test_isfinite() | ||
| 191 | + self._test_unique_dim() | ||
| 192 | + self._test_npu_prompt_flash_attention() | ||
| 193 | + | ||
| 194 | + | ||
| 92 | if __name__ == '__main__': | 195 | if __name__ == '__main__': |
| 93 | run_tests() | 196 | run_tests() |