已合并
feat: Add batch_norm_reduce API documentation and enhance unit tests. #5251
fgd_dragon创建于 6月24日
feat: Add batch_norm_reduce API documentation and enhance unit tests. #5251
已合并
共 7 个文件变更+143-18
| @@ -506,6 +506,43 @@ tensor([ 14.7733, -30.1218, -1.3647, ..., -16.6840, 7.1518, 8.4872], | |||
| 506 | ) | 506 | ) |
| 507 | 507 | ||
| 508 | 508 | ||
| 509 | +_add_torch_npu_docstr( | ||
| 510 | + "batch_norm_reduce", | ||
| 511 | + """ | ||
| 512 | +功能描述: | ||
| 513 | +该接口用于按BatchNorm通道维对输入Tensor进行规约，计算通道维以外维度上的元素和sum与平方和square_sum。 | ||
| 514 | + | ||
| 515 | +接口原型: | ||
| 516 | +torch_npu.batch_norm_reduce(input, eps) -> (Tensor, Tensor) | ||
| 517 | + | ||
| 518 | +参数说明: | ||
| 519 | +input (Tensor): 必选参数，表示待规约的输入Tensor。支持4D Tensor，shape为(N, C, H, W)。数据类型支持float32、float16、bfloat16。 | ||
| 520 | +eps (float): 必选参数，保留参数。该参数不参与sum和square_sum的计算，建议按BatchNorm常用配置传入1e-5。 | ||
| 521 | + | ||
| 522 | +输出说明: | ||
| 523 | +sum (Tensor): 输入Tensor按通道维规约后的元素和，为1D float32 Tensor，长度为通道维长度。输入为4D Tensor且shape为(N, C, H, W)时，输出shape为(C,)。 | ||
| 524 | +square_sum (Tensor): 输入Tensor按通道维规约后的平方和，为1D float32 Tensor，shape同sum。 | ||
| 525 | + | ||
| 526 | +支持的型号: | ||
| 527 | +Atlas A3 训练系列产品/Atlas A3 推理系列产品 | ||
| 528 | +Atlas A2 训练系列产品/Atlas A2 推理系列产品 | ||
| 529 | + | ||
| 530 | +调用示例: | ||
| 531 | +import torch | ||
| 532 | +import torch_npu | ||
| 533 | + | ||
| 534 | +input = torch.randn(2, 3, 12, 12, dtype=torch.float32, device="npu") | ||
| 535 | + | ||
| 536 | +sum_out, square_sum_out = torch_npu.batch_norm_reduce(input, 1e-5) | ||
| 537 | + | ||
| 538 | +print("sum_out.shape:", sum_out.shape) | ||
| 539 | +print("sum_out.dtype:", sum_out.dtype) | ||
| 540 | +print("square_sum_out.shape:", square_sum_out.shape) | ||
| 541 | +print("square_sum_out.dtype:", square_sum_out.dtype) | ||
| 542 | +""" | ||
| 543 | +) | ||
| 544 | + | ||
| 545 | + | ||
| 509 | _add_torch_npu_docstr( | 546 | _add_torch_npu_docstr( |
| 510 | "npu_block_sparse_attention", | 547 | "npu_block_sparse_attention", |
| 511 | """ | 548 | """ |
| @@ -0,0 +1,54 @@ | |||
| 1 | +# torch_npu.batch_norm_reduce | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| --- | :---: | | ||
| 7 | +| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ | | ||
| 8 | +| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ | | ||
| 9 | + | ||
| 10 | +## 功能说明 | ||
| 11 | + | ||
| 12 | +- API功能：该接口用于按BatchNorm通道维对输入Tensor进行规约，计算通道维以外维度上的元素和 `sum` 与平方和 `square_sum` 。 | ||
| 13 | + | ||
| 14 | +- 计算公式：对于4D Tensor输入 `(N, C, H, W)`，计算C轴以外的N、H、W轴求和与平方和。 | ||
| 15 | + | ||
| 16 | + $$ | ||
| 17 | + sum_i = \sum_{n=0}^{N-1} \sum_{h=0}^{H-1} \sum_{w=0}^{W-1} input_{(n,i,h,w)} | ||
| 18 | + $$ | ||
| 19 | + | ||
| 20 | + $$ | ||
| 21 | + square\_sum_i = \sum_{n=0}^{N-1} \sum_{h=0}^{H-1} \sum_{w=0}^{W-1} input_{(n,i,h,w)}^2 | ||
| 22 | + $$ | ||
| 23 | + | ||
| 24 | +## 函数原型 | ||
| 25 | + | ||
| 26 | +```python | ||
| 27 | +torch_npu.batch_norm_reduce(input, eps) -> (Tensor, Tensor) | ||
| 28 | +``` | ||
| 29 | + | ||
| 30 | +## 参数说明 | ||
| 31 | + | ||
| 32 | +- **input** (`Tensor`)：必选参数，表示待规约的输入Tensor。支持4D Tensor，shape为 `(N, C, H, W)` 。数据类型支持 `float32`、`float16`、`bfloat16`。 | ||
| 33 | +- **eps** (`float`)：必选参数，保留参数。该参数不参与 `sum` 和 `square_sum` 的计算，建议按BatchNorm常用配置传入 `1e-5`。 | ||
| 34 | + | ||
| 35 | +## 返回值说明 | ||
| 36 | + | ||
| 37 | +- **sum** (`Tensor`)：输入Tensor按通道维规约后的元素和，为1D `float32` Tensor，长度为通道维长度。输入为4D Tensor且shape为 `(N, C, H, W)` 时，输出shape为 `(C,)` 。 | ||
| 38 | +- **square_sum** (`Tensor`)：输入Tensor按通道维规约后的平方和，为1D `float32` Tensor，shape同 `sum` 。 | ||
| 39 | + | ||
| 40 | +## 调用示例 | ||
| 41 | + | ||
| 42 | +```python | ||
| 43 | +import torch | ||
| 44 | +import torch_npu | ||
| 45 | + | ||
| 46 | +input = torch.randn(2, 3, 12, 12, dtype=torch.float32, device="npu") | ||
| 47 | + | ||
| 48 | +sum_out, square_sum_out = torch_npu.batch_norm_reduce(input, 1e-5) | ||
| 49 | + | ||
| 50 | +print("sum_out.shape:", sum_out.shape) | ||
| 51 | +print("sum_out.dtype:", sum_out.dtype) | ||
| 52 | +print("square_sum_out.shape:", square_sum_out.shape) | ||
| 53 | +print("square_sum_out.dtype:", square_sum_out.dtype) | ||
| 54 | +``` | ||
| @@ -797,6 +797,11 @@ | |||
| 797 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p><a name="p1281815913362"></a><a name="p1281815913362"></a>该接口实现了Lightning Indexer组件warmup阶段训练的反向梯度计算，并融合了Loss的计算。</p> | 797 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p><a name="p1281815913362"></a><a name="p1281815913362"></a>该接口实现了Lightning Indexer组件warmup阶段训练的反向梯度计算，并融合了Loss的计算。</p> |
| 798 | </td> | 798 | </td> |
| 799 | </tr> | 799 | </tr> |
| 800 | +<tr id="batch_norm_reduce"><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p id="p_batch_norm_reduce"><a name="p_batch_norm_reduce"></a><a name="p_batch_norm_reduce"></a><a href="torch_npu-batch_norm_reduce.md">torch_npu.batch_norm_reduce</a></p> | ||
| 801 | +</td> | ||
| 802 | +<td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p_batch_norm_reduce_desc"><a name="p_batch_norm_reduce_desc"></a><a name="p_batch_norm_reduce_desc"></a><span>对输入Tensor按BatchNorm通道维进行规约，返回元素和与平方和。</span></p> | ||
| 803 | +</td> | ||
| 804 | +</tr> | ||
| 800 | <tr><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p><a href="torch_npu-npu_add_rms_norm.md">torch_npu.npu_add_rms_norm</a></p> | 805 | <tr><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p><a href="torch_npu-npu_add_rms_norm.md">torch_npu.npu_add_rms_norm</a></p> |
| 801 | </td> | 806 | </td> |
| 802 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p>将Add计算与RMSNorm归一化融合，常用于大模型中将残差连接后的张量进行归一化处理。</p> | 807 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p>将Add计算与RMSNorm归一化融合，常用于大模型中将残差连接后的张量进行归一化处理。</p> |
| @@ -6903,8 +6903,8 @@ def convolution_meta(input_tensor, weight, bias, stride, padding, dilation, tran | |||
| 6903 | 6903 | ||
| 6904 | 6904 | ||
| 6905 | def batch_norm_reduce_meta(self, eps): | 6905 | def batch_norm_reduce_meta(self, eps): |
| 6906 | - out_sum = torch.empty(self.size(1), dtype=self.dtype, device='meta') | 6906 | + out_sum = torch.empty(self.size(1), dtype=torch.float32, device='meta') |
| 6907 | - out_square_sum = torch.empty(self.size(1), dtype=self.dtype, device='meta') | 6907 | + out_square_sum = torch.empty(self.size(1), dtype=torch.float32, device='meta') |
| 6908 | return (out_sum, out_square_sum) | 6908 | return (out_sum, out_square_sum) |
| 6909 | 6909 | ||
| 6910 | 6910 | ||
| @@ -5605,13 +5605,13 @@ class TestNpuConvolutionTranspose(TestCase): | |||
| 5605 | class TestBatchNormReduce(TestCase): | 5605 | class TestBatchNormReduce(TestCase): |
| 5606 | def test_batch_norm_reduce(self): | 5606 | def test_batch_norm_reduce(self): |
| 5607 | with FakeTensorMode(): | 5607 | with FakeTensorMode(): |
| 5608 | - input_fake_tensor = torch.randn(2, 3, 12, 12, device="npu", requires_grad=True) | 5608 | + input_fake_tensor = torch.randn(2, 3, 12, 12, dtype=torch.float16, device="npu", requires_grad=True) |
| 5609 | eps = 1e-5 | 5609 | eps = 1e-5 |
| 5610 | output_fake_tensor1, output_fake_tensor2 = torch_npu.batch_norm_reduce(input_fake_tensor, eps) | 5610 | output_fake_tensor1, output_fake_tensor2 = torch_npu.batch_norm_reduce(input_fake_tensor, eps) |
| 5611 | self.assertEqual(input_fake_tensor.shape[1], output_fake_tensor1.shape[0]) | 5611 | self.assertEqual(input_fake_tensor.shape[1], output_fake_tensor1.shape[0]) |
| 5612 | - self.assertEqual(input_fake_tensor.dtype, output_fake_tensor1.dtype) | 5612 | + self.assertEqual(torch.float32, output_fake_tensor1.dtype) |
| 5613 | self.assertEqual(input_fake_tensor.shape[1], output_fake_tensor2.shape[0]) | 5613 | self.assertEqual(input_fake_tensor.shape[1], output_fake_tensor2.shape[0]) |
| 5614 | - self.assertEqual(input_fake_tensor.dtype, output_fake_tensor2.dtype) | 5614 | + self.assertEqual(torch.float32, output_fake_tensor2.dtype) |
| 5615 | 5615 | ||
| 5616 | 5616 | ||
| 5617 | class TestMatmul(TestCase): | 5617 | class TestMatmul(TestCase): |
| @@ -3,33 +3,61 @@ import numpy as np | |||
| 3 | 3 | ||
| 4 | import torch_npu | 4 | import torch_npu |
| 5 | from torch_npu.testing.testcase import TestCase, run_tests | 5 | from torch_npu.testing.testcase import TestCase, run_tests |
| 6 | -from torch_npu.testing.common_utils import create_common_tensor | 6 | +from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices |
| 7 | 7 | ||
| 8 | 8 | ||
| 9 | class TestBatchNormReduce(TestCase): | 9 | class TestBatchNormReduce(TestCase): |
| 10 | def cuda_op_exec(self, input_data): | 10 | def cuda_op_exec(self, input_data): |
| 11 | - cpu_sum = torch.sum(input_data, dim=[0, 2, 3]) | 11 | + input_float = input_data.to(torch.float32) |
| 12 | - cpu_square_sum = torch.sum(input_data * input_data, dim=[0, 2, 3]) | 12 | + cpu_sum = torch.sum(input_float, dim=[0, 2, 3]) |
| 13 | + cpu_square_sum = torch.sum(input_float * input_float, dim=[0, 2, 3]) | ||
| 13 | return cpu_sum.numpy(), cpu_square_sum.numpy() | 14 | return cpu_sum.numpy(), cpu_square_sum.numpy() |
| 14 | 15 | ||
| 15 | def npu_op_exec(self, *args): | 16 | def npu_op_exec(self, *args): |
| 16 | - npu_sum, npu_square_sum = torch_npu.batch_norm_reduce(*args) | 17 | + return torch_npu.batch_norm_reduce(*args) |
| 17 | - out_sum = npu_sum.cpu().numpy() | ||
| 18 | - out_square_sum = npu_square_sum.cpu().numpy() | ||
| 19 | - return out_sum, out_square_sum | ||
| 20 | 18 | ||
| 21 | - def test_batch_norm_reduce(self): | 19 | + def assert_batch_norm_reduce_result(self, cpu_input, npu_input, eps=1e-5, rtol=1e-3, atol=1e-3): |
| 20 | + cpu_sum, cpu_square_sum = self.cuda_op_exec(cpu_input) | ||
| 21 | + npu_sum, npu_square_sum = self.npu_op_exec(npu_input, eps) | ||
| 22 | + | ||
| 23 | + self.assertEqual((cpu_input.shape[1],), tuple(npu_sum.shape)) | ||
| 24 | + self.assertEqual((cpu_input.shape[1],), tuple(npu_square_sum.shape)) | ||
| 25 | + self.assertEqual(torch.float32, npu_sum.dtype) | ||
| 26 | + self.assertEqual(torch.float32, npu_square_sum.dtype) | ||
| 27 | + self.assertRtolEqual(cpu_sum, npu_sum.cpu().numpy(), rtol, atol) | ||
| 28 | + self.assertRtolEqual(cpu_square_sum, npu_square_sum.cpu().numpy(), rtol, atol) | ||
| 29 | + | ||
| 30 | + def test_batch_norm_reduce_normal_cases(self): | ||
| 22 | np.random.seed(1234) | 31 | np.random.seed(1234) |
| 23 | shape_format = [ | 32 | shape_format = [ |
| 24 | [[np.float32, -1, [2, 3, 12, 12]], 1e-5], | 33 | [[np.float32, -1, [2, 3, 12, 12]], 1e-5], |
| 34 | + [[np.float32, -1, [1, 1, 2, 2]], 1e-5], | ||
| 35 | + [[np.float32, -1, [3, 7, 1, 2]], 1e-5], | ||
| 36 | + [[np.float16, -1, [2, 3, 12, 12]], 1e-5], | ||
| 25 | ] | 37 | ] |
| 26 | for item in shape_format: | 38 | for item in shape_format: |
| 27 | - cpu_input1, npu_inputfp32 = create_common_tensor(item[0], 1, 10) | 39 | + cpu_input, npu_input = create_common_tensor(item[0], 1, 10) |
| 28 | - cpu_output = self.cuda_op_exec(cpu_input1) | 40 | + self.assert_batch_norm_reduce_result(cpu_input, npu_input, item[-1]) |
| 29 | - npu_outputfp32 = self.npu_op_exec(npu_inputfp32, item[-1]) | ||
| 30 | 41 | ||
| 31 | - self.assertRtolEqual(cpu_output[0], npu_outputfp32[0]) | 42 | + @SupportedDevices(['Ascend910B']) |
| 32 | - self.assertRtolEqual(cpu_output[1], npu_outputfp32[1], 1e-2) | 43 | + def test_batch_norm_reduce_bfloat16(self): |
| 44 | + torch.manual_seed(1234) | ||
| 45 | + cpu_input = torch.randn(2, 3, 12, 12, dtype=torch.float32).to(torch.bfloat16) | ||
| 46 | + npu_input = cpu_input.npu() | ||
| 47 | + self.assert_batch_norm_reduce_result(cpu_input, npu_input, rtol=4e-3, atol=4e-3) | ||
| 48 | + | ||
| 49 | + def test_batch_norm_reduce_eps_no_effect(self): | ||
| 50 | + _, npu_input = create_common_tensor([np.float32, -1, [2, 3, 12, 12]], 1, 10) | ||
| 51 | + npu_sum1, npu_square_sum1 = self.npu_op_exec(npu_input, 1e-5) | ||
| 52 | + npu_sum2, npu_square_sum2 = self.npu_op_exec(npu_input, 1.0) | ||
| 53 | + | ||
| 54 | + self.assertRtolEqual(npu_sum1.cpu().numpy(), npu_sum2.cpu().numpy()) | ||
| 55 | + self.assertRtolEqual(npu_square_sum1.cpu().numpy(), npu_square_sum2.cpu().numpy()) | ||
| 56 | + | ||
| 57 | + def test_batch_norm_reduce_invalid_dim(self): | ||
| 58 | + npu_input = torch.randn(3, dtype=torch.float32).npu() | ||
| 59 | + with self.assertRaisesRegex(RuntimeError, "dim input tensor|must more than 1"): | ||
| 60 | + torch_npu.batch_norm_reduce(npu_input, 1e-5) | ||
| 33 | 61 | ||
| 34 | 62 | ||
| 35 | if __name__ == "__main__": | 63 | if __name__ == "__main__": |
menu位置和list位置要保持一致