已合并
[fix] 无反向算子实现的API ,求梯度抛异常 #5447
L1919_snow创建于 7月10日
[fix] 无反向算子实现的API ,求梯度抛异常 #5447
已合并
共 2 个文件变更+77-0
| @@ -0,0 +1,17 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | + | ||
| 4 | + | ||
| 5 | + // 以下算子没有反向实现。为它们注册 PyTorch 内置的 autogradNotImplementedFallback： | ||
| 6 | + // 前向正常执行，只有在对依赖这些算子的张量调用 .backward() 时，才会抛出 | ||
| 7 | + // "derivative for 'npu::xxx' is not implemented" 异常。 | ||
| 8 | + // 该 fallback 是 boxed 内核，对任意算子 schema 通用，无需匹配各自签名。 | ||
| 9 | + TORCH_LIBRARY_IMPL(npu, AutogradPrivateUse1, m) { | ||
| 10 | + m.impl("npu_add_rms_norm", torch::autograd::autogradNotImplementedFallback()); | ||
| 11 | + m.impl("npu_interleave_rope", torch::autograd::autogradNotImplementedFallback()); | ||
| 12 | + m.impl("npu_moe_gating_top_k_softmax", torch::autograd::autogradNotImplementedFallback()); | ||
| 13 | + // npu_apply_rotary_pos_emb 仅在 v2.7+ 暴露，低版本未注册该算子，需加版本保护。 | ||
| 14 | + | ||
| 15 | + m.impl("npu_apply_rotary_pos_emb", torch::autograd::autogradNotImplementedFallback()); | ||
| 16 | + | ||
| 17 | + } | ||
| @@ -0,0 +1,60 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch_npu | ||
| 3 | + | ||
| 4 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 5 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +# 以下 4 个算子没有反向实现，已在 AutogradPrivateUse1 上注册 | ||
| 9 | +# autogradNotImplementedFallback：前向正常，对输出调用 .backward() 时应抛出 | ||
| 10 | +# RuntimeError: derivative for 'npu::xxx' is not implemented。 | ||
| 11 | +# 仅在 A5 / Ascend950 上验证，其它机型由 SupportedDevices 自动跳过。 | ||
| 12 | +class TestAutogradErrorFallback(TestCase): | ||
| 13 | + | ||
| 14 | + def _leaf(self, *shape, dtype=torch.float16): | ||
| 15 | + x = torch.randn(*shape, dtype=dtype).npu() | ||
| 16 | + x.requires_grad_(True) | ||
| 17 | + return x | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + def test_npu_add_rms_norm_backward_raises(self): | ||
| 21 | + x1 = self._leaf(4, 16) | ||
| 22 | + x2 = self._leaf(4, 16) | ||
| 23 | + gamma = self._leaf(16) | ||
| 24 | + out = torch_npu.npu_add_rms_norm(x1, x2, gamma, 1e-6)[0] | ||
| 25 | + self.assertRaisesRegex( | ||
| 26 | + RuntimeError, "not implemented", | ||
| 27 | + lambda: out.float().sum().backward()) | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + def test_npu_apply_rotary_pos_emb_backward_raises(self): | ||
| 31 | + q = self._leaf(1, 4, 2, 16) | ||
| 32 | + k = self._leaf(1, 4, 2, 16) | ||
| 33 | + cos = torch.randn(1, 4, 1, 16, dtype=torch.float16).npu() | ||
| 34 | + sin = torch.randn(1, 4, 1, 16, dtype=torch.float16).npu() | ||
| 35 | + out = torch_npu.npu_apply_rotary_pos_emb(q, k, cos, sin, "BSND", "half")[0] | ||
| 36 | + self.assertRaisesRegex( | ||
| 37 | + RuntimeError, "not implemented", | ||
| 38 | + lambda: out.float().sum().backward()) | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + def test_npu_interleave_rope_backward_raises(self): | ||
| 42 | + x = self._leaf(1, 4, 2, 16) | ||
| 43 | + cos = torch.randn(1, 4, 1, 16, dtype=torch.float16).npu() | ||
| 44 | + sin = torch.randn(1, 4, 1, 16, dtype=torch.float16).npu() | ||
| 45 | + out = torch_npu.npu_interleave_rope(x, cos, sin) | ||
| 46 | + self.assertRaisesRegex( | ||
| 47 | + RuntimeError, "not implemented", | ||
| 48 | + lambda: out.float().sum().backward()) | ||
| 49 | + | ||
| 50 | + | ||
| 51 | + def test_npu_moe_gating_top_k_softmax_backward_raises(self): | ||
| 52 | + x = self._leaf(8, 16) | ||
| 53 | + out = torch_npu.npu_moe_gating_top_k_softmax(x, None, 2)[0] | ||
| 54 | + self.assertRaisesRegex( | ||
| 55 | + RuntimeError, "not implemented", | ||
| 56 | + lambda: out.float().sum().backward()) | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +if __name__ == "__main__": | ||
| 60 | + run_tests() | ||