Star847
1330
代码
代码
Issues1512
Pull Requests2236
流水线
项目讨论
Wiki
分析
Star847
1330
已合并
[test][v2.12.0] add NPU validation cases for torch._functorch APIs (issues #2684 #2685 #2687 #2688) #41386
ggg_0963创建于 7月12日
[test][v2.12.0] add NPU validation cases for torch._functorch APIs (issues #2684 #2685 #2687 #2688) #41386
已合并
ggg_0963创建于 7月12日
从ggg_0963/pytorchtest-functorch-aot-config-vmap-v2.12.0合入到Ascend/pytorchv2.12.0
共 4 个文件变更+618-0
Atest/functorch/test_aot_joint_with_descriptors_api.py+96-0
@@ -0,0 +1,96 @@
1+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
2+#
3+# Licensed under the BSD 3-Clause License (the "License");
4+# you may not use this file except in compliance with the License.
5+# You may obtain a copy of the License at
6+#
7+# https://opensource.org/licenses/BSD-3-Clause
8+#
9+# Unless required by applicable law or agreed to in writing, software
10+# distributed under the License is distributed on an "AS IS" BASIS,
11+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+# See the License for the specific language governing permissions and
13+# limitations under the License.
14+ 
15+"""
16+Minimal direct-API validation for the two AOT joint-with-descriptors entrypoints on NPU:
17+1. torch._functorch.aot_autograd.aot_export_joint_with_descriptors
18+2. torch._functorch.aot_autograd.aot_compile_joint_with_descriptors
19+ 
20+This file deliberately avoids reusing the upstream
21+test/functorch/test_aot_joint_with_descriptors.py scaffolding (named module
22+classes, assertExpectedInline FX graph text comparison, decomposition_table,
23+full backward correctness). Instead it asserts only the observable contract
24+of the API pair on NPU against an eager reference built from nn.Sequential:
25+the export call returns a JointWithDescriptors exposing graph_module and
26+_aot_state; the compile call returns a callable whose forward result matches
27+the eager module.
28+ 
29+Invocation contract: the callable returned by aot_compile_joint_with_descriptors
30+flattens (params, inputs) into positional args via fx_pytree, so it must be
31+called as compiled(*params, *inputs), matching the upstream release/2.9+
32+test convention `parallel_model_fn(*dict(model.named_parameters()).values(), *inputs)`.
33+"""
34+from contextlib import ExitStack
35+ 
36+import torch
37+ 
38+from torch._functorch.aot_autograd import (
39+ aot_compile_joint_with_descriptors,
40+ aot_export_joint_with_descriptors,
41+)
42+from torch.testing._internal.common_utils import run_tests, TestCase
43+ 
44+device_type = (
45+ acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
46+)
47+ 
48+ 
49+def _build_export_target():
50+ """An eager nn.Sequential reference plus matching inputs on the active device."""
51+ model = torch.nn.Sequential(torch.nn.Linear(2, 1)).to(device_type)
52+ inputs = (torch.randn(3, 2, device=device_type),)
53+ return model, inputs
54+ 
55+ 
56+class TestAOTJointWithDescriptorsNPU(TestCase):
57+ def test_export_returns_joint_with_descriptors(self):
58+ """aot_export_joint_with_descriptors returns a JointWithDescriptors
59+ exposing graph_module and _aot_state on NPU.
60+ """
61+ model, inputs = _build_export_target()
62+ with ExitStack() as stack:
63+ exported = aot_export_joint_with_descriptors(stack, model, inputs)
64+ self.assertIsNotNone(exported)
65+ self.assertIsNotNone(exported.graph_module)
66+ self.assertIsNotNone(exported._aot_state)
67+ 
68+ def test_export_preserves_npu_device(self):
69+ """aot_export_joint_with_descriptors leaves input tensors and
70+ module parameters on the NPU device.
71+ """
72+ model, inputs = _build_export_target()
73+ with ExitStack() as stack:
74+ aot_export_joint_with_descriptors(stack, model, inputs)
75+ self.assertEqual(inputs[0].device.type, device_type)
76+ for p in model.parameters():
77+ self.assertEqual(p.device.type, device_type)
78+ 
79+ def test_compile_runs_and_matches_eager(self):
80+ """aot_compile_joint_with_descriptors runs end-to-end on NPU and
81+ the compiled forward matches the eager module output.
82+ """
83+ model, inputs = _build_export_target()
84+ with ExitStack() as stack:
85+ exported = aot_export_joint_with_descriptors(stack, model, inputs)
86+ compiled = aot_compile_joint_with_descriptors(exported)
87+ self.assertTrue(callable(compiled))
88+ expected = model(*inputs)
89+ actual = compiled(*dict(model.named_parameters()).values(), *inputs)
90+ self.assertEqual(actual.shape, expected.shape)
91+ self.assertEqual(actual.device, expected.device)
92+ self.assertEqual(actual, expected, rtol=1e-3, atol=1e-3)
93+ 
94+ 
95+if __name__ == "__main__":
96+ run_tests()
Atest/functorch/test_functorch_config_api.py+102-0
@@ -0,0 +1,102 @@
1+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
2+#
3+# Licensed under the BSD 3-Clause License (the "License");
4+# you may not use this file except in compliance with the License.
5+# You may obtain a copy of the License at
6+#
7+# https://opensource.org/licenses/BSD-3-Clause
8+#
9+# Unless required by applicable law or agreed to in writing, software
10+# distributed under the License is distributed on an "AS IS" BASIS,
11+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+# See the License for the specific language governing permissions and
13+# limitations under the License.
14+ 
15+"""
16+Add validation cases for torch._functorch.config.patch API on NPU:
17+1. PyTorch community lacks sufficient and direct API validations for this API, so this file is added.
18+2. This file validates torch._functorch.config.patch (extendable).
19+"""
20+import torch
21+import torch._functorch.config as config
22+from torch.testing._internal.common_utils import run_tests, TestCase
atomgit-bot
atomgit-botatomgit-bot7月12日

🔴 Critical

文件 test/functorch/test_functorch_config_api.py 使用 device_type = "npu" 并通过 .to(device_type) 将张量移至 NPU 设备，但未导入 torch_npu。在本仓库中，torch_npu 负责注册 NPU 设备；若不导入，torch.randn(...).to("npu") 将抛出 RuntimeError: Device type 'npu' is not supported。

对比同目录下已有测试文件 test/functorch/test_eager_transforms.py（第 25-26 行）明确导入了 import torch_npu 和 import torch_npu.testing。

触发条件：在 NPU 环境下运行此测试文件时，所有测试方法均会因 NPU 设备未注册而失败。

建议：在 import torch 之后添加 import torch_npu 和 import torch_npu.testing，与同目录下 test_eager_transforms.py 保持一致。

改动建议
22
+ import torch
23
+ import torch_npu
24
+ import torch_npu.testing
22
- from torch.testing._internal.common_utils import run_tests, TestCase
25
+ import torch._functorch.config as config
应用建议
likedislike
群
群青世界
7月24日 评论：
23+ 
24+device_type = (
25+ acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
26+)
27+ 
28+ 
29+class TestFunctorchConfigPatch(TestCase):
30+ def test_basic_patch(self):
31+ """Test basic config.patch context manager"""
32+ original_value = config.debug_assert
33+ with config.patch("debug_assert", not original_value):
34+ self.assertEqual(config.debug_assert, not original_value)
35+ self.assertEqual(config.debug_assert, original_value)
36+ 
37+ def test_patch_dict(self):
38+ """Test config.patch with dict argument"""
39+ original_debug = config.debug_assert
40+ original_cse = config.cse
41+ patches = {
42+ "debug_assert": not original_debug,
43+ "cse": not original_cse,
44+ }
45+ with config.patch(patches):
46+ self.assertEqual(config.debug_assert, not original_debug)
47+ self.assertEqual(config.cse, not original_cse)
48+ self.assertEqual(config.debug_assert, original_debug)
49+ self.assertEqual(config.cse, original_cse)
50+ 
51+ def test_patch_restore_after_exception(self):
52+ """Test that config is restored even after an exception"""
53+ original_value = config.debug_assert
54+ with self.assertRaises(RuntimeError):
55+ with config.patch("debug_assert", not original_value):
56+ self.assertEqual(config.debug_assert, not original_value)
57+ raise RuntimeError("test exception")
58+ self.assertEqual(config.debug_assert, original_value)
59+ 
60+ def test_patch_nested(self):
61+ """Test nested config.patch contexts"""
62+ original_value = config.debug_assert
63+ # First level
64+ with config.patch("debug_assert", True):
65+ self.assertTrue(config.debug_assert)
66+ # Second level (nested)
67+ with config.patch("debug_assert", False):
68+ self.assertFalse(config.debug_assert)
69+ # Back to first level
70+ self.assertTrue(config.debug_assert)
71+ # Back to original
72+ self.assertEqual(config.debug_assert, original_value)
73+ 
74+ def test_patch_with_tensor_device(self):
75+ """Test that config.patch works correctly with NPU tensors in the context"""
76+ x = torch.randn(3, 4).to(device_type)
77+ original_value = config.debug_assert
78+ with config.patch("debug_assert", not original_value):
79+ # Verify NPU tensor operations still work
80+ y = x.mm(x.T)
81+ self.assertEqual(y.device.type, device_type)
82+ self.assertEqual(config.debug_assert, not original_value)
83+ self.assertEqual(config.debug_assert, original_value)
84+ 
85+ def test_patch_invalid_key(self):
86+ """Test patch with invalid key raises AttributeError."""
87+ with self.assertRaises(AttributeError):
88+ with config.patch("non_existent_key", 42):
89+ pass
90+ 
91+ def test_patch_invalid_dict_key(self):
92+ """Test patch dict with invalid key among valid keys raises AttributeError and does not change valid keys."""
93+ original_value = config.debug_assert
94+ with self.assertRaises(AttributeError):
95+ with config.patch({"debug_assert": True, "invalid_key": 42}):
96+ pass
97+ # Verify valid key is unchanged after the failed patch
98+ self.assertEqual(config.debug_assert, original_value)
99+ 
100+ 
101+if __name__ == "__main__":
102+ run_tests()
Atest/functorch/test_vmap_add_batch_dim_api.py+184-0
@@ -0,0 +1,184 @@
1+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
2+#
3+# Licensed under the BSD 3-Clause License (the "License");
4+# you may not use this file except in compliance with the License.
5+# You may obtain a copy of the License at
6+#
7+# https://opensource.org/licenses/BSD-3-Clause
8+#
9+# Unless required by applicable law or agreed to in writing, software
10+# distributed under the License is distributed on an "AS IS" BASIS,
11+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+# See the License for the specific language governing permissions and
13+# limitations under the License.
14+ 
15+"""
16+Add validation cases for torch._functorch.vmap._add_batch_dim API on NPU:
17+1. PyTorch community lacks sufficient and direct API validations for this API, so this file is added.
18+2. This file validates torch._functorch.vmap._add_batch_dim (extendable).
19+"""
20+import torch
21+from torch._functorch.vmap import _add_batch_dim
22+from torch.testing._internal.common_utils import run_tests, TestCase
atomgit-bot
atomgit-botatomgit-bot7月12日

🔴 Critical

文件 test/functorch/test_vmap_add_batch_dim_api.py 使用 device_type = "npu" 并通过 .to(device_type) 将张量移至 NPU 设备，但未导入 torch_npu。在本仓库中，torch_npu 负责注册 NPU 设备；若不导入，torch.randn(...).to("npu") 将抛出 RuntimeError: Device type 'npu' is not supported。

对比同目录下已有测试文件 test/functorch/test_eager_transforms.py（第 25-26 行）明确导入了 import torch_npu 和 import torch_npu.testing。

触发条件：在 NPU 环境下运行此测试文件时，所有测试方法均会因 NPU 设备未注册而失败。

建议：在 import torch 之后添加 import torch_npu 和 import torch_npu.testing，与同目录下 test_eager_transforms.py 保持一致。

改动建议
22
+ import torch
23
+ import torch_npu
24
+ import torch_npu.testing
22
- from torch.testing._internal.common_utils import run_tests, TestCase
25
+ from torch._functorch.vmap import _add_batch_dim
应用建议
likedislike
群
群青世界
7月24日 评论：
23+ 
24+device_type = (
25+ acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
26+)
27+ 
28+ 
29+class TestVmapAddBatchDim(TestCase):
30+ def test_add_batch_dim_basic(self):
31+ """Test basic _add_batch_dim functionality on NPU"""
32+ x = torch.randn(3, 4).to(device_type)
33+ vmap_level = 0
34+ batched = _add_batch_dim(x, 0, vmap_level)
35+ self.assertIsNotNone(batched)
36+ self.assertIsInstance(batched, torch.Tensor)
37+ self.assertEqual(batched.device.type, device_type)
38+ self.assertEqual(batched.shape, (4,))
39+ 
40+ def test_add_batch_dim_with_vmap(self):
41+ """Test _add_batch_dim works correctly with vmap on NPU"""
42+ x = torch.randn(2, 3).to(device_type)
43+ y = torch.randn(2, 3).to(device_type)
44+ 
45+ def dot_row(a, b):
46+ return (a * b).sum(dim=-1)
47+ 
48+ result = torch.vmap(dot_row)(x, y)
49+ expected = dot_row(x, y)
50+ self.assertEqual(result, expected)
51+ self.assertEqual(result.device.type, device_type)
52+ 
53+ def test_add_batch_dim_nested_vmap(self):
54+ """Test nested vmap with _add_batch_dim on NPU"""
55+ x = torch.randn(2, 3, 4).to(device_type)
56+ y = torch.randn(2, 3, 4).to(device_type)
57+ 
58+ def matmul_row(a, b):
59+ return (a * b).sum(dim=-1)
60+ 
61+ result = torch.vmap(torch.vmap(matmul_row))(x, y)
62+ self.assertEqual(result.shape, (2, 3))
63+ self.assertEqual(result.device.type, device_type)
64+ 
65+ def test_add_batch_dim_with_model(self):
66+ """Test _add_batch_dim with a simple model on NPU"""
67+ model = torch.nn.Linear(4, 2).to(device_type)
68+ x = torch.randn(3, 4).to(device_type)
69+ 
70+ result = torch.vmap(lambda x: model(x))(x)
71+ expected = model(x)
72+ self.assertEqual(result, expected)
73+ self.assertEqual(result.device.type, device_type)
74+ 
75+ def test_add_batch_dim_in_dims(self):
76+ """Test _add_batch_dim with different in_dims on NPU"""
77+ x = torch.randn(3, 4, 5).to(device_type)
78+ 
79+ def identity(x):
80+ return x
81+ 
82+ # Test in_dims=0 (default)
83+ result0 = torch.vmap(identity, in_dims=0)(x)
84+ self.assertEqual(result0.shape, (3, 4, 5))
85+ 
86+ # Test in_dims=1
87+ result1 = torch.vmap(identity, in_dims=1)(x)
88+ self.assertEqual(result1.shape, (4, 3, 5))
89+ 
90+ # Test in_dims=-1
91+ result_neg1 = torch.vmap(identity, in_dims=-1)(x)
92+ self.assertEqual(result_neg1.shape, (5, 3, 4))
93+ 
94+ def test_add_batch_dim_out_dims(self):
95+ """Test _add_batch_dim with different out_dims on NPU"""
96+ x = torch.randn(3, 4).to(device_type)
97+ 
98+ def identity(x):
99+ return x
100+ 
101+ # Test out_dims=0 (default)
102+ result0 = torch.vmap(identity, out_dims=0)(x)
103+ self.assertEqual(result0.shape, (3, 4))
104+ 
105+ # Test out_dims=1
106+ result1 = torch.vmap(identity, out_dims=1)(x)
107+ self.assertEqual(result1.shape, (4, 3))
108+ 
109+ def test_add_batch_dim_with_grad(self):
110+ """Test _add_batch_dim works with gradient computation on NPU"""
111+ x = torch.randn(3, 3, device=device_type, requires_grad=True)
112+ w = torch.randn(3, 3, device=device_type, requires_grad=True)
113+ 
114+ def fn(x, w):
115+ return (x * w).sum(dim=-1)
116+ 
117+ result = torch.vmap(fn)(x, w)
118+ loss = result.sum()
119+ loss.backward()
120+ self.assertIsNotNone(x.grad)
121+ self.assertIsNotNone(w.grad)
122+ self.assertEqual(x.grad.device.type, device_type)
123+ self.assertEqual(w.grad.device.type, device_type)
124+ 
125+ def test_add_batch_dim_direct_3d_batch_dim_0(self):
126+ """Direct API call: 3D tensor with batch_dim=0 returns shape=(4,5) (inner value sliced at dim 0)."""
127+ x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type)
128+ vmap_level = 0
129+ batched = _add_batch_dim(x, 0, vmap_level)
130+ self.assertIsNotNone(batched)
131+ self.assertEqual(batched.shape, (4, 5))
132+ self.assertEqual(batched.device.type, device_type)
133+ self.assertEqual(batched.dtype, torch.float32)
134+ 
135+ def test_add_batch_dim_direct_3d_batch_dim_1(self):
136+ """Direct API call: 3D tensor with batch_dim=1 returns shape=(3,5)."""
137+ x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type)
138+ vmap_level = 0
139+ batched = _add_batch_dim(x, 1, vmap_level)
140+ self.assertIsNotNone(batched)
141+ self.assertEqual(batched.shape, (3, 5))
142+ self.assertEqual(batched.device.type, device_type)
143+ 
144+ def test_add_batch_dim_direct_3d_batch_dim_2(self):
145+ """Direct API call: 3D tensor with batch_dim=2 returns shape=(3,4)."""
146+ x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type)
147+ vmap_level = 0
148+ batched = _add_batch_dim(x, 2, vmap_level)
149+ self.assertIsNotNone(batched)
150+ self.assertEqual(batched.shape, (3, 4))
151+ self.assertEqual(batched.device.type, device_type)
152+ 
153+ def test_add_batch_dim_direct_preserves_dtype_and_device(self):
154+ """Direct API call preserves dtype and device."""
155+ x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32).to(device_type)
156+ vmap_level = 0
157+ batched = _add_batch_dim(x, 0, vmap_level)
158+ self.assertEqual(batched.dtype, torch.float32)
159+ self.assertEqual(batched.device.type, device_type)
160+ self.assertEqual(batched.shape, (2,))
161+ 
162+ def test_add_batch_dim_invalid_batch_dim(self):
163+ """Direct API call with out-of-range batch_dim wraps (x.ndim % batch_dim)."""
164+ x = torch.randn(3, 4).to(device_type)
165+ # batch_dim=5 on 2D tensor acts as batch_dim=1 (5 % 2 ≡ 1)
166+ batched = _add_batch_dim(x, 5, 0)
167+ self.assertEqual(batched.shape, (3,))
168+ self.assertEqual(batched.device.type, device_type)
169+ 
170+ def test_add_batch_dim_multiple_levels(self):
171+ """Test different vmap_levels produce independent batch dimensions."""
172+ from torch._functorch.vmap import _remove_batch_dim
173+ x = torch.randn(3, 4).to(device_type)
174+ # Level 0
175+ b_l0 = _add_batch_dim(x, 0, 0)
176+ # Level 1 - independent from level 0
177+ b_l1 = _add_batch_dim(x, 0, 1)
178+ # Both levels show the same shape on the surface
179+ self.assertEqual(b_l0.shape, (4,))
180+ self.assertEqual(b_l1.device.type, device_type)
181+ 
182+ 
183+if __name__ == "__main__":
184+ run_tests()
Atest_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch+236-0
@@ -0,0 +1,236 @@
1+diff --git a/test/functorch/test_aot_joint_with_descriptors.py b/test/functorch/test_aot_joint_with_descriptors.py
2+index 13f2318..0dfba56 100644
3+--- a/test/functorch/test_aot_joint_with_descriptors.py
4++++ b/test/functorch/test_aot_joint_with_descriptors.py
5+@@ -39,13 +39,17 @@ from torch._functorch.aot_autograd import (
6+ from torch._guards import tracing, TracingContext
7+ from torch.nn.attention.flex_attention import create_block_mask, flex_attention
8+ from torch.testing._internal.common_utils import (
9+- requires_cuda,
10+ run_tests,
11+ skipIfCrossRef,
12+ TestCase,
13+ )
14+
15+
16++device_type = (
17++ acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
18++)
19++
20++
21+ def graph_capture(model, inputs, with_export):
22+ gm = model
23+ tracing_context = None
24+@@ -76,8 +80,8 @@ class TestAOTJointWithDescriptors(TestCase):
25+ def forward(self, x):
26+ return self.linear(x)
27+
28+- model = SimpleLinear()
29+- inputs = (torch.randn(4, 3),)
30++ model = SimpleLinear().to(device_type)
31++ inputs = (torch.randn(4, 3, device=device_type),)
32+
33+ with ExitStack() as stack:
34+ # Export joint with descriptors
35+@@ -152,9 +156,9 @@ class inner_f(torch.nn.Module):
36+ x = self.bn(x)
37+ return torch.relu(x)
38+
39+- model = ConvBN()
40++ model = ConvBN().to(device_type)
41+ model.train() # Important for batch norm
42+- inputs = (torch.randn(2, 1, 4, 4),)
43++ inputs = (torch.randn(2, 1, 4, 4, device=device_type),)
44+
45+ with ExitStack() as stack:
46+ # Export joint with descriptors
47+@@ -317,9 +321,9 @@ class inner_f(torch.nn.Module):
48+ def forward(self, x, *, scale):
49+ return self.linear(x) * scale
50+
51+- model = ModuleWithKwargs()
52+- inputs = (torch.randn(4, 3),)
53+- kwargs = {"scale": torch.tensor(2.0)}
54++ model = ModuleWithKwargs().to(device_type)
55++ inputs = (torch.randn(4, 3, device=device_type),)
56++ kwargs = {"scale": torch.tensor(2.0, device=device_type)}
57+
58+ gm = dynamo_graph_capture_for_export(model)(*inputs, **kwargs)
59+
60+@@ -408,8 +412,8 @@ class inner_f(torch.nn.Module):
61+ out2 = self.linear2(x)
62+ return out1, out2
63+
64+- model = MultiOutputModule()
65+- inputs = (torch.randn(4, 3),)
66++ model = MultiOutputModule().to(device_type)
67++ inputs = (torch.randn(4, 3, device=device_type),)
68+
69+ with ExitStack() as stack:
70+ # Export joint with descriptors
71+@@ -505,8 +509,8 @@ class inner_f(torch.nn.Module):
72+ def forward(self, x):
73+ return self.linear(x)
74+
75+- model = SimpleModule()
76+- inputs = (torch.randn(4, 3),)
77++ model = SimpleModule().to(device_type)
78++ inputs = (torch.randn(4, 3, device=device_type),)
79+
80+ with ExitStack() as stack:
81+ # Export joint with descriptors
82+@@ -554,8 +558,8 @@ class inner_f(torch.nn.Module):
83+ def forward(self, x):
84+ return self.linear(x)
85+
86+- model = SimpleLinear()
87+- inputs = (torch.randn(4, 3),)
88++ model = SimpleLinear().to(device_type)
89++ inputs = (torch.randn(4, 3, device=device_type),)
90+
91+ with ExitStack() as stack:
92+ joint_with_descriptors = aot_export_joint_with_descriptors(
93+@@ -630,9 +634,9 @@ class inner_f(torch.nn.Module):
94+ x = self.bn(x)
95+ return torch.relu(x)
96+
97+- model = ConvBN()
98++ model = ConvBN().to(device_type)
99+ model.train() # Important for batch norm
100+- inputs = (torch.randn(2, 1, 4, 4),)
101++ inputs = (torch.randn(2, 1, 4, 4, device=device_type),)
102+
103+ with ExitStack() as stack:
104+ joint_with_descriptors = aot_export_joint_with_descriptors(
105+@@ -693,8 +697,8 @@ class inner_f(torch.nn.Module):
106+ out2 = self.linear2(x)
107+ return out1, out2
108+
109+- model = MultiOutputModule()
110+- inputs = (torch.randn(4, 3),)
111++ model = MultiOutputModule().to(device_type)
112++ inputs = (torch.randn(4, 3, device=device_type),)
113+
114+ with ExitStack() as stack:
115+ joint_with_descriptors = aot_export_joint_with_descriptors(
116+@@ -738,8 +742,8 @@ class inner_f(torch.nn.Module):
117+ def forward(self, x):
118+ return self.linear(x)
119+
120+- model = SimpleModule()
121+- inputs = (torch.randn(4, 3),)
122++ model = SimpleModule().to(device_type)
123++ inputs = (torch.randn(4, 3, device=device_type),)
124+
125+ with ExitStack() as stack:
126+ joint_with_descriptors = aot_export_joint_with_descriptors(
127+@@ -783,8 +787,8 @@ class inner_f(torch.nn.Module):
128+ def forward(self, x):
129+ return self.linear(x)
130+
131+- model = SimpleModule()
132+- inputs = (torch.randn(4, 3),)
133++ model = SimpleModule().to(device_type)
134++ inputs = (torch.randn(4, 3, device=device_type),)
135+
136+ with ExitStack() as stack:
137+ joint_with_descriptors = aot_export_joint_with_descriptors(
138+@@ -809,8 +813,8 @@ class inner_f(torch.nn.Module):
139+ y = self.linear(x)
140+ return y - 1
141+
142+- inputs = (torch.randn(4, 3),)
143+- model = SimpleLinear()
144++ inputs = (torch.randn(4, 3, device=device_type),)
145++ model = SimpleLinear().to(device_type)
146+
147+ for with_export in [True, False]:
148+ graph_module = graph_capture(model, inputs, with_export)
149+@@ -828,7 +832,6 @@ class inner_f(torch.nn.Module):
150+ ('call_function', 't_3', {'pp_stage': 0})""",
151+ )
152+
atomgit-botatomgit-bot
atomgit-botatomgit-bot7月12日
已过期

🟠 High Priority

在 patch 文件第 153 行，对 FlexAttentionModule 的处理为：

此处先 .to(device_type) 将模型移至 NPU，随后又 .to(device) 将模型移回原设备。然而同一测试中 xq、xk、xv 的输入张量仅做了 .to(device_type) 移至 NPU。模型与输入张量设备不一致将导致运行时错误（如 RuntimeError: Expected all tensors to be on the same device）。

触发条件：当 device_type 为 "npu" 且 device 为 "cuda" 或 "cpu" 时，模型最终在原设备上而输入在 NPU 上，前向计算将因设备不匹配而失败。

建议：将 FlexAttentionModule().to(device_type).to(device) 改为 FlexAttentionModule().to(device_type)，避免双重设备移动导致模型与输入设备不一致。

改动建议
155
-
155
+ + model = FlexAttentionModule().to(device_type)
应用建议
likedislike
群
群青世界
7月24日 评论：
atomgit-botatomgit-bot7月12日
已过期

🟡 Medium Priority

patch 第 153 行将 model = FlexAttentionModule().to(device) 改为 model = FlexAttentionModule().to(device_type).to(device)。该链式调用先将模型移至 device_type（如 "npu"），随后又立即移至 device（原始代码中的设备，如 "cuda"），最终模型停留在 device 上。而同一测试中，tensor（xq、xk、xv）通过 .to(device_type) 被移至 device_type（patch 第 142/146/150 行）。这导致模型在 device 上而输入 tensor 在 device_type 上，运行时将因设备不匹配而报错。

变更行：patch 第 153 行 +model = FlexAttentionModule().to(device_type).to(device) → model 最终在 device 上，tensor 在 device_type 上 → 设备不匹配 RuntimeError → 应改为 model = FlexAttentionModule().to(device_type)，去掉末尾的 .to(device)。

建议：将 model = FlexAttentionModule().to(device_type).to(device) 改为 model = FlexAttentionModule().to(device_type)，确保 model 与 tensor 在同一设备上。

改动建议
155
-
155
+ + model = FlexAttentionModule().to(device_type)
应用建议
likedislike
群
群青世界
7月24日 评论：
153+- @requires_cuda
154+ def test_preserve_annotate_flex_attention(self):
155+ def score_mod(score, b, h, m, n):
156+ return score
157+@@ -844,7 +847,7 @@ class inner_f(torch.nn.Module):
158+ b = 24
159+ batch_size = 2
160+ seqlen = a * b
161+- device = "cuda"
162++ device = device_type
163+
164+ # Create seq_idx tensor - maps each position to a document/sequence ID
165+ # Example: Split sequence into 2 documents for each batch
166+@@ -951,8 +954,8 @@ class inner_f(torch.nn.Module):
167+ y = example_function(y)
168+ return y - 1
169+
170+- inputs = (torch.randn(4, 3),)
171+- model = SimpleLinear()
172++ inputs = (torch.randn(4, 3, device=device_type),)
173++ model = SimpleLinear().to(device_type)
174+
175+ for with_export in [True, False]:
176+ graph_module = graph_capture(model, inputs, with_export)
177+@@ -995,8 +998,8 @@ class inner_f(torch.nn.Module):
178+ def forward(self, x, y):
179+ return foo(x, y)
180+
181+- model = CustomOpModule()
182+- inputs = (torch.randn(4, 3), torch.randn(4, 3))
183++ model = CustomOpModule().to(device_type)
184++ inputs = (torch.randn(4, 3, device=device_type), torch.randn(4, 3, device=device_type))
185+
186+ gm = graph_capture(model, inputs, with_export=True)
187+
188+@@ -1028,7 +1031,7 @@ class inner_f(torch.nn.Module):
189+ def __init__(self):
190+ super().__init__()
191+ self.input_shape = (5, 3)
192+- self.permuted_indices = torch.tensor([2, 0, 3, 1])
193++ self.permuted_indices = torch.tensor([2, 0, 3, 1], device=device_type)
194+
195+ def forward(self, x):
196+ with fx_traceback.annotate({"pp_stage": 0}):
197+@@ -1037,7 +1040,7 @@ class inner_f(torch.nn.Module):
198+ )
199+ return routed_output.cos()
200+
201+- inputs = (torch.randn(4, 3, requires_grad=True),)
202++ inputs = (torch.randn(4, 3, device=device_type, requires_grad=True),)
203+ model = Module()
204+
205+ graph_module = graph_capture(model, inputs, True)
206+@@ -1078,8 +1081,8 @@ class inner_f(torch.nn.Module):
207+ def forward(self, x):
208+ return self.linear(x)
209+
210+- model = SimpleLinear()
211+- inputs = (torch.randn(4, 3),)
212++ model = SimpleLinear().to(device_type)
213++ inputs = (torch.randn(4, 3, device=device_type),)
214+ gm = dynamo_graph_capture_for_export(model)(*inputs)
215+ fake_mode = gm.meta.get("fake_mode", None)
216+
217+@@ -1105,8 +1108,8 @@ class inner_f(torch.nn.Module):
218+ with fx_traceback.annotate({"test": 1}):
219+ return self.linear(x) - self.linear2(x)
220+
221+- model = SimpleLinear()
222+- inputs = (torch.randn(4, 3, requires_grad=True),)
223++ model = SimpleLinear().to(device_type)
224++ inputs = (torch.randn(4, 3, device=device_type, requires_grad=True),)
225+ graph_module = graph_capture(model, inputs, True)
226+ add_nodes = graph_module.graph.find_nodes(
227+ op="call_function", target=torch.ops.aten.add.Tensor
228+@@ -1161,7 +1164,7 @@ class inner_f(torch.nn.Module):
229+ z = self.bar(x)
230+ return z - 1
231+
232+- inputs = (torch.randn(4, 3, requires_grad=True),)
233++ inputs = (torch.randn(4, 3, device=device_type, requires_grad=True),)
234+ model = MyMod()
235+
236+ # invoke_subgraph doesn't seem to work with with_export=False, no subgraph created