已合并
fix: defer ACLGraph update until after checkpoint #39072
luochao60创建于 6月23日
fix: defer ACLGraph update until after checkpoint #39072
已合并
共 3 个文件变更+23-35
| @@ -540,9 +540,9 @@ class TestACLGraphUpdatePlanCompile(TestUtils): | |||
| 540 | original_update = graph_tree.update_aclgraph_records_for_graph | 540 | original_update = graph_tree.update_aclgraph_records_for_graph |
| 541 | seen_plans = [] | 541 | seen_plans = [] |
| 542 | 542 | ||
| 543 | - def collect_plan(plan, graph, inputs): | 543 | + def collect_plan(update_input, graph): |
| 544 | - seen_plans.append(plan) | 544 | + seen_plans.append(update_input) |
| 545 | - return original_update(plan, graph, inputs) | 545 | + return original_update(update_input, graph) |
| 546 | 546 | ||
| 547 | try: | 547 | try: |
| 548 | config.triton.cudagraphs = True | 548 | config.triton.cudagraphs = True |
| @@ -578,13 +578,9 @@ class TestACLGraphUpdatePlanCompile(TestUtils): | |||
| 578 | torch._dynamo.reset() | 578 | torch._dynamo.reset() |
| 579 | 579 | ||
| 580 | self.assertTrue(seen_plans) | 580 | self.assertTrue(seen_plans) |
| 581 | - self.assertTrue(any(plan for plan in seen_plans)) | 581 | + self.assertTrue(any(update_input for update_input in seen_plans)) |
| 582 | - plan = next(plan for plan in seen_plans if plan) | 582 | + update_input = next(update_input for update_input in seen_plans if update_input) |
| 583 | - self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score.default") | 583 | + self.assertEqual(update_input[0]["actual_seq_lengths"], [37]) |
| 584 | - self.assertEqual( | ||
| 585 | - plan[0]["updates"]["actual_seq_lengths"], | ||
| 586 | - {"kind": "list", "items": [{"kind": "constant", "value": 37}]}, | ||
| 587 | - ) | ||
| 588 | 584 | ||
| 589 | def test_npugraphify_keeps_aclgraph_update_plan_on_callable_attribute(self): | 585 | def test_npugraphify_keeps_aclgraph_update_plan_on_callable_attribute(self): |
| 590 | import torch_npu.npu._graph_tree as graph_tree | 586 | import torch_npu.npu._graph_tree as graph_tree |
| @@ -637,9 +633,9 @@ class TestACLGraphUpdatePlanCompile(TestUtils): | |||
| 637 | original_update = graph_tree.update_aclgraph_records_for_graph | 633 | original_update = graph_tree.update_aclgraph_records_for_graph |
| 638 | seen_plans = [] | 634 | seen_plans = [] |
| 639 | 635 | ||
| 640 | - def collect_plan(plan, graph, inputs): | 636 | + def collect_plan(update_input, graph): |
| 641 | - seen_plans.append(plan) | 637 | + seen_plans.append(update_input) |
| 642 | - return original_update(plan, graph, inputs) | 638 | + return original_update(update_input, graph) |
| 643 | 639 | ||
| 644 | try: | 640 | try: |
| 645 | config.triton.cudagraphs = True | 641 | config.triton.cudagraphs = True |
| @@ -674,13 +670,9 @@ class TestACLGraphUpdatePlanCompile(TestUtils): | |||
| 674 | config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts | 670 | config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts |
| 675 | torch._dynamo.reset() | 671 | torch._dynamo.reset() |
| 676 | 672 | ||
| 677 | - self.assertTrue(any(plan for plan in seen_plans)) | 673 | + self.assertTrue(any(update_input for update_input in seen_plans)) |
| 678 | - plan = next(plan for plan in seen_plans if plan) | 674 | + update_input = next(update_input for update_input in seen_plans if update_input) |
| 679 | - self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score_v2.default") | 675 | + self.assertEqual(update_input[0]["actual_seq_qlen"], [1]) |
| 680 | - self.assertEqual( | ||
| 681 | - plan[0]["updates"]["actual_seq_qlen"], | ||
| 682 | - {"kind": "list", "items": [{"kind": "constant", "value": 1}]}, | ||
| 683 | - ) | ||
| 684 | 676 | ||
| 685 | 677 | ||
| 686 | if __name__ == "__main__": | 678 | if __name__ == "__main__": |
| @@ -142,16 +142,15 @@ def validate_aclgraph_update_plan_for_graph( | |||
| 142 | 142 | ||
| 143 | 143 | ||
| 144 | def update_aclgraph_records_for_graph( | 144 | def update_aclgraph_records_for_graph( |
| 145 | - plan: Sequence[Dict[str, Any]], | 145 | + cpu_update_input: Sequence[Dict[str, Any]], |
| 146 | graph: Any, | 146 | graph: Any, |
| 147 | - new_inputs: Sequence[Any], | ||
| 148 | ) -> bool: | 147 | ) -> bool: |
 | |||
| 149 | if graph is None or not graph.auto_dispatch_capture: | 148 | if graph is None or not graph.auto_dispatch_capture: |
| 150 | return False | 149 | return False |
| 151 | - if not plan: | 150 | + if not cpu_update_input: |
| 152 | return False | 151 | return False |
| 153 | 152 | ||
| 154 | - graph.update(resolve_aclgraph_update_plan(plan, new_inputs)) | 153 | + graph.update(cpu_update_input) |
| 155 | return True | 154 | return True |
| 156 | 155 | ||
| 157 | 156 | ||
| @@ -106,6 +106,7 @@ from torch_npu._C import ( | |||
| 106 | _set_cached_tensors_enabled as _set_cached_tensors_enabled) | 106 | _set_cached_tensors_enabled as _set_cached_tensors_enabled) |
| 107 | from torch_npu.npu._aclgraph_update_plan.resolver import ( | 107 | from torch_npu.npu._aclgraph_update_plan.resolver import ( |
| 108 | ACLGRAPH_UPDATE_PLAN_GLOBAL, | 108 | ACLGRAPH_UPDATE_PLAN_GLOBAL, |
| 109 | + resolve_aclgraph_update_plan, | ||
| 109 | update_aclgraph_records_for_graph, | 110 | update_aclgraph_records_for_graph, |
| 110 | validate_aclgraph_update_plan_for_graph, | 111 | validate_aclgraph_update_plan_for_graph, |
| 111 | ) | 112 | ) |
| @@ -1069,9 +1070,8 @@ class NPUGraphNode: | |||
| 1069 | log.debug("NPUGRAPH-TREE Node Run node=%s", self.id) | 1070 | log.debug("NPUGRAPH-TREE Node Run node=%s", self.id) |
| 1070 | self.check_static_inputs_are_stable(new_inputs) | 1071 | self.check_static_inputs_are_stable(new_inputs) |
| 1071 | aclgraph_update_submitted = update_aclgraph_records_for_graph( | 1072 | aclgraph_update_submitted = update_aclgraph_records_for_graph( |
| 1072 | - self.aclgraph_update_plan, | 1073 | + resolve_aclgraph_update_plan(self.aclgraph_update_plan, new_inputs), |
| 1073 | self.graph, | 1074 | self.graph, |
| 1074 | - new_inputs, | ||
| 1075 | ) | 1075 | ) |
| 1076 | self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs) | 1076 | self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs) |
| 1077 | 1077 | ||
| @@ -1244,7 +1244,10 @@ class NPUGraphNode: | |||
| 1244 | 1244 | ||
| 1245 | check_memory_pool(self.device, self.npu_graphs_pool, memory) | 1245 | check_memory_pool(self.device, self.npu_graphs_pool, memory) |
| 1246 | 1246 | ||
| 1247 | - aclgraph_update_inputs = list(inputs) | 1247 | + aclgraph_cpu_update_input = resolve_aclgraph_update_plan( |
| 1248 | + self.aclgraph_update_plan, | ||
| 1249 | + inputs, | ||
| 1250 | + ) | ||
| 1248 | 1251 | ||
| 1249 | with preserve_rng_state(), torch.npu.device( | 1252 | with preserve_rng_state(), torch.npu.device( |
| 1250 | self.device | 1253 | self.device |
| @@ -1257,14 +1260,6 @@ class NPUGraphNode: | |||
| 1257 | ), get_history_recording(): | 1260 | ), get_history_recording(): |
| 1258 | static_outputs = model(inputs) | 1261 | static_outputs = model(inputs) |
| 1259 | 1262 | ||
| 1260 | - validate_aclgraph_update_plan_for_graph(self.aclgraph_update_plan, self.graph) | ||
| 1261 | - update_aclgraph_records_for_graph( | ||
| 1262 | - self.aclgraph_update_plan, | ||
| 1263 | - self.graph, | ||
| 1264 | - aclgraph_update_inputs, | ||
| 1265 | - ) | ||
| 1266 | - aclgraph_update_inputs.clear() | ||
| 1267 | - | ||
| 1268 | # running model should reclaim memory | 1263 | # running model should reclaim memory |
| 1269 | if not len(inputs) == 0: | 1264 | if not len(inputs) == 0: |
| 1270 | raise RuntimeError("check len(inputs) == 0 fail") | 1265 | raise RuntimeError("check len(inputs) == 0 fail") |
| @@ -1272,6 +1267,8 @@ class NPUGraphNode: | |||
| 1272 | static_outputs = (static_outputs,) | 1267 | static_outputs = (static_outputs,) |
| 1273 | 1268 | ||
| 1274 | self._add_first_outputs(static_outputs, static_input_persistent_storage_ptrs) | 1269 | self._add_first_outputs(static_outputs, static_input_persistent_storage_ptrs) |
| 1270 | + validate_aclgraph_update_plan_for_graph(self.aclgraph_update_plan, self.graph) | ||
| 1271 | + update_aclgraph_records_for_graph(aclgraph_cpu_update_input, self.graph) | ||
| 1275 | 1272 | ||
| 1276 | log.debug("NPUGRAPH-TREE Node Record node=%s recorded: outputs=%d, " | 1273 | log.debug("NPUGRAPH-TREE Node Record node=%s recorded: outputs=%d, " |
| 1277 | "non_static_inputs=%d, static_input_idxs=%d", | 1274 | "non_static_inputs=%d, static_input_idxs=%d", |
🟡 Medium Priority
changed line:
resolver.py中update_aclgraph_records_for_graph函数签名从(plan, graph, new_inputs)3 参数改为(cpu_update_input, graph)2 参数。affected behavior/contract: 测试文件
test/_inductor/test_aclgraph_update_plan_compile.py在两个测试函数中（line 543 和 line 640）通过 monkey-patch 替换update_aclgraph_records_for_graph为collect_plan(plan, graph, inputs)包装函数，该包装函数期望 3 个参数。包装函数内部也以 3 参数调用保存的original_update。failure mode: 当生产代码以新的 2 参数形式
update_aclgraph_records_for_graph(cpu_update_input, graph)调用时，实际执行的是collect_plan，但只传入 2 个参数，导致TypeError: collect_plan() missing 1 required positional argument: 'inputs'。即使包装函数被调用成功，其内部的original_update(plan, graph, inputs)也会因传入 3 参数给新 2 参数函数而抛出TypeError。测试将直接崩溃而非提供有意义的失败信息。suggested fix: 更新测试文件中的两处 mock 包装函数，使其签名为
collect_plan(cpu_update_input, graph)以匹配新接口。同时更新内部对original_update(cpu_update_input, graph)的调用。建议：更新 test/_inductor/test_aclgraph_update_plan_compile.py 中两处 collect_plan 包装函数签名（line 543 和 line 640），将参数从 (plan, graph, inputs) 改为 (cpu_update_input, graph)，并相应调整函数体内对 original_update 的调用和 seen_plans 的记录逻辑。