已合并
add npu_gemma_rms_norm #2208
liuyun_nj创建于 2025年3月4日
add npu_gemma_rms_norm #2208
已合并
从refs/pull/2208/head合入到master
共 3 个文件变更+100-0
| @@ -7095,6 +7095,9 @@ custom: | |||
| 7095 | op_api: all_version | 7095 | op_api: all_version |
| 7096 | exposed: all_version | 7096 | exposed: all_version |
| 7097 | 7097 | ||
| 7098 | + - func: npu_gemma_rms_norm(Tensor self, Tensor gamma, float epsilon=1e-06) -> (Tensor, Tensor) | ||
| 7099 | + op_api: all_version | ||
| 7100 | + | ||
| 7098 | - func: npu_rms_norm_backward(Tensor dy, Tensor self, Tensor gamma, Tensor rstd) -> (Tensor, Tensor) | 7101 | - func: npu_rms_norm_backward(Tensor dy, Tensor self, Tensor gamma, Tensor rstd) -> (Tensor, Tensor) |
| 7099 | acl_op: all_version | 7102 | acl_op: all_version |
| 7100 | op_api: all_version | 7103 | op_api: all_version |
| @@ -0,0 +1,35 @@ | |||
| 1 | +// Copyright (c) 2025 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +namespace op_api { | ||
| 21 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 22 | + | ||
| 23 | +std::tuple<at::Tensor, at::Tensor> npu_gemma_rms_norm( | ||
| 24 | + const at::Tensor& self, | ||
| 25 | + const at::Tensor& gamma, | ||
| 26 | + double epsilon) | ||
| 27 | +{ | ||
| 28 | + auto output_size = op_infer::rms_norm_npu_output_size(self, gamma); | ||
| 29 | + at::Tensor y = npu_preparation::apply_tensor_with_format(output_size[0], self.options(), ACL_FORMAT_ND); | ||
| 30 | + at::Tensor rstd = npu_preparation::apply_tensor_with_format(output_size[1], self.options().dtype(at::kFloat), ACL_FORMAT_ND); | ||
| 31 | + EXEC_NPU_CMD(aclnnGemmaRmsNorm, self, gamma, epsilon, y, rstd); | ||
| 32 | + return std::tuple<at::Tensor, at::Tensor>(y, rstd); | ||
| 33 | +} | ||
| 34 | + | ||
| 35 | +} | ||
| @@ -0,0 +1,62 @@ | |||
| 1 | +import unittest | ||
| 2 | +import numpy as np | ||
| 3 | +import torch_npu | ||
| 4 | +import torch | ||
| 5 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 6 | +from torch_npu.testing.common_utils import get_npu_device, SupportedDevices | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestNPUGemmaRmsNorm(TestCase): | ||
| 10 | + | ||
| 11 | + def supported_op_exec(self, x, gamma): | ||
| 12 | + x_fp32 = np.array(x, dtype=np.float32) | ||
| 13 | + gamma_fp32 = np.array(gamma, dtype=np.float32) | ||
| 14 | + | ||
| 15 | + variance = np.mean(np.power(x_fp32, 2), axis=-1, keepdims=True) | ||
| 16 | + epsilon = 1e-6 | ||
| 17 | + std = np.sqrt(variance + epsilon) | ||
| 18 | + rstd = 1 / std | ||
| 19 | + result_mid = x_fp32 * rstd | ||
| 20 | + gamma_mid = gamma_fp32 + 1 | ||
| 21 | + result_fp32 = result_mid * gamma_mid | ||
| 22 | + | ||
| 23 | + result = np.array(result_fp32, dtype=x.dtype) | ||
| 24 | + | ||
| 25 | + return result, rstd | ||
| 26 | + | ||
| 27 | + def custom_op_exec(self, x, gamma): | ||
| 28 | + y, rstd = torch_npu.npu_gemma_rms_norm(x, gamma) | ||
| 29 | + return y.cpu().numpy(), rstd.cpu().numpy() | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + def test_gemma_rms_norm(self, device="npu"): | ||
| 34 | + if device is None: | ||
| 35 | + device = get_npu_device() | ||
| 36 | + cpu_input0 = np.random.uniform(0, 100, [256, 512]).astype(np.float32) | ||
| 37 | + cpu_input1 = np.random.uniform(0, 100, [512]).astype(np.float32) | ||
| 38 | + npu_input0 = torch.from_numpy(cpu_input0).to(device) | ||
| 39 | + npu_input1 = torch.from_numpy(cpu_input1).to(device) | ||
| 40 | + | ||
| 41 | + supported_output0, supported_output1 = self.supported_op_exec(cpu_input0, cpu_input1) | ||
| 42 | + custom_output0, custom_output1 = self.custom_op_exec(npu_input0, npu_input1) | ||
| 43 | + self.assertRtolEqual(supported_output0, custom_output0) | ||
| 44 | + self.assertRtolEqual(supported_output1, custom_output1) | ||
| 45 | + | ||
| 46 | + | ||
| 47 | + | ||
| 48 | + def test_gemma_rms_norm_fp16(self, device="npu"): | ||
| 49 | + if device is None: | ||
| 50 | + device = get_npu_device() | ||
| 51 | + cpu_input0 = np.random.uniform(0, 100, [256, 512]).astype(np.float16) | ||
| 52 | + cpu_input1 = np.random.uniform(0, 100, [512]).astype(np.float16) | ||
| 53 | + npu_input0 = torch.from_numpy(cpu_input0).to(device) | ||
| 54 | + npu_input1 = torch.from_numpy(cpu_input1).to(device) | ||
| 55 | + | ||
| 56 | + supported_output0, supported_output1 = self.supported_op_exec(cpu_input0, cpu_input1) | ||
| 57 | + custom_output0, custom_output1 = self.custom_op_exec(npu_input0, npu_input1) | ||
| 58 | + self.assertRtolEqual(supported_output0, custom_output0) | ||
| 59 | + self.assertRtolEqual(supported_output1, custom_output1) | ||
| 60 | + | ||
| 61 | +if __name__ == "__main__": | ||
| 62 | + run_tests() | ||