已合并
torch.nn.functional.layer_norm calls aclnnLayerNorm only on A5 #5473
chengpeng25创建于 7月14日
torch.nn.functional.layer_norm calls aclnnLayerNorm only on A5 #5473
已合并
共 2 个文件变更+38-1
| @@ -94,7 +94,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> native_layer_norm(const at::Tenso | |||
| 94 | } | 94 | } |
| 95 | // call HostAPI function | 95 | // call HostAPI function |
| 96 | static auto layer_sc = at_npu::native::env::CheckCompatibleImpl(); | 96 | static auto layer_sc = at_npu::native::env::CheckCompatibleImpl(); |
| 97 | - if (!layer_sc) { | 97 | + if (!layer_sc || c10_npu::IsAclnnOnly()) { |
| 98 | EXEC_NPU_CMD(aclnnLayerNorm, input, normalized_shape, input_weight, input_bias, eps, output, mean_out, rstd_out); | 98 | EXEC_NPU_CMD(aclnnLayerNorm, input, normalized_shape, input_weight, input_bias, eps, output, mean_out, rstd_out); |
| 99 | } else { | 99 | } else { |
| 100 | EXEC_NPU_CMD(aclnnFastLayerNorm, input, normalized_shape, input_weight, input_bias, eps, output, mean_out, rstd_out); | 100 | EXEC_NPU_CMD(aclnnFastLayerNorm, input, normalized_shape, input_weight, input_bias, eps, output, mean_out, rstd_out); |
| @@ -113,6 +113,43 @@ class TestLayerNorm(TestCase): | |||
| 113 | self.assertTrue("Given normalized_shape=[2, 3, 3], expected input with shape [*, 2, 3, 3], but got input of size[10, 5]" in str(exception)) | 113 | self.assertTrue("Given normalized_shape=[2, 3, 3], expected input with shape [*, 2, 3, 3], but got input of size[10, 5]" in str(exception)) |
| 114 | 114 | ||
| 115 | 115 | ||
| 116 | + | ||
| 117 | + def test_layer_norm_compatible_impl_switch(self): | ||
| 118 | + shape_format = [ | ||
| 119 | + [np.float32, 0, (64, 10)], | ||
| 120 | + [np.float32, 0, (256, 2048, 7, 7)], | ||
| 121 | + [np.float16, 0, (10, 128)], | ||
| 122 | + [np.float16, 0, (46, 16)], | ||
| 123 | + ] | ||
| 124 | + | ||
| 125 | + try: | ||
| 126 | + for item in shape_format: | ||
| 127 | + cpu_input, npu_input = create_common_tensor(item, 1, 100) | ||
| 128 | + normalized_shape = cpu_input.size()[1:] | ||
| 129 | + | ||
| 130 | + cpu_output = torch.nn.functional.layer_norm( | ||
| 131 | + cpu_input.float(), normalized_shape) | ||
| 132 | + if item[0] == np.float16: | ||
| 133 | + cpu_output = cpu_output.to(torch.float16) | ||
| 134 | + | ||
| 135 | + torch_npu.npu.use_compatible_impl(False) | ||
| 136 | + npu_out_no_compat = torch.nn.functional.layer_norm( | ||
| 137 | + npu_input, normalized_shape) | ||
| 138 | + self.assertRtolEqual( | ||
| 139 | + | ||
| 140 | + cpu_output.detach().numpy(), | ||
| 141 | + npu_out_no_compat.cpu().detach().numpy()) | ||
| 142 | + | ||
| 143 | + torch_npu.npu.use_compatible_impl(True) | ||
| 144 | + npu_out_with_compat = torch.nn.functional.layer_norm( | ||
| 145 | + npu_input, normalized_shape) | ||
| 146 | + self.assertRtolEqual( | ||
| 147 | + cpu_output.detach().numpy(), | ||
| 148 | + npu_out_with_compat.cpu().detach().numpy()) | ||
| 149 | + finally: | ||
| 150 | + torch_npu.npu.use_compatible_impl(True) | ||
| 151 | + | ||
| 152 | + | ||
| 116 | if __name__ == "__main__": | 153 | if __name__ == "__main__": |
| 117 | torch_npu.npu.use_compatible_impl(True) | 154 | torch_npu.npu.use_compatible_impl(True) |
| 118 | run_tests() | 155 | run_tests() |