已合并
【feature】remove elu and elubackward in decomp #37502
hbhu_bin创建于 6月3日
【feature】remove elu and elubackward in decomp #37502
已合并
共 2 个文件变更+54-0
| @@ -0,0 +1,52 @@ | |||
| 1 | +# Owner(s): ["module: tests"] | ||
| 2 | +"""End-to-end: elu/elu_backward are numerically correct through the DVM | ||
| 3 | +backend's native aclnn fallback (elu is excluded from decomposition; the policy | ||
| 4 | +assertions live in test_dvm_decomp.py). | ||
| 5 | + | ||
| 6 | +The DVM backend is pinned via torch.compile(options={"npu_backend": "dvm"}). | ||
| 7 | +We deliberately do NOT import torch_npu._inductor at module scope: importing it | ||
| 8 | +loads a backend at import time, which would turn the first torch.compile into a | ||
| 9 | +mid-process backend switch (default -> dvm). Plain ``import torch_npu`` does not | ||
| 10 | +load _inductor. | ||
| 11 | +""" | ||
| 12 | +import unittest | ||
| 13 | + | ||
| 14 | +import torch | ||
| 15 | +import torch_npu | ||
| 16 | +from torch.testing._internal.common_utils import ( | ||
| 17 | + instantiate_parametrized_tests, | ||
| 18 | + parametrize, | ||
| 19 | + run_tests, | ||
| 20 | + TestCase, | ||
| 21 | +) | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +class TestDvmEluNative(TestCase): | ||
| 26 | + | ||
| 27 | + def test_elu_forward_backward_matches_eager(self, dtype): | ||
| 28 | + tol = 1e-3 if dtype == torch.float32 else 4e-3 | ||
| 29 | + # build in fp32 then cast (npu normal kernel has no bf16 support) | ||
| 30 | + ref = torch.randn(256, 4096, device="npu").to(dtype) | ||
| 31 | + x_e = ref.detach().clone().requires_grad_(True) | ||
| 32 | + x_c = ref.detach().clone().requires_grad_(True) | ||
| 33 | + | ||
| 34 | + def fn(t): | ||
| 35 | + return torch.nn.functional.elu(t) | ||
| 36 | + | ||
| 37 | + out_e = fn(x_e) | ||
| 38 | + out_e.float().sum().backward() | ||
| 39 | + | ||
| 40 | + compiled = torch.compile( | ||
| 41 | + fn, backend="inductor", options={"npu_backend": "dvm"} | ||
| 42 | + ) | ||
| 43 | + out_c = compiled(x_c) | ||
| 44 | + out_c.float().sum().backward() | ||
| 45 | + | ||
| 46 | + self.assertEqual(out_e, out_c, atol=tol, rtol=tol) | ||
| 47 | + self.assertEqual(x_e.grad, x_c.grad, atol=tol, rtol=tol) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +instantiate_parametrized_tests(TestDvmEluNative) | ||
| 51 | +if __name__ == "__main__": | ||
| 52 | + run_tests() | ||
| @@ -24,6 +24,8 @@ decomps_to_exclude_npu = [ | |||
| 24 | aten.embedding_dense_backward, | 24 | aten.embedding_dense_backward, |
| 25 | aten.gelu.default, | 25 | aten.gelu.default, |
| 26 | aten.gelu_backward.default, | 26 | aten.gelu_backward.default, |
| 27 | + aten.elu.default, | ||
| 28 | + aten.elu_backward.default, | ||
| 27 | aten.grid_sampler_2d, | 29 | aten.grid_sampler_2d, |
| 28 | aten.grid_sampler_2d_backward, | 30 | aten.grid_sampler_2d_backward, |
| 29 | aten.linalg_vector_norm, | 31 | aten.linalg_vector_norm, |