已合并
Add support of Execution Trace Observer object to torch_npu profiler (based on acl_prof). #44622
ilya_a创建于 18 天前
Add support of Execution Trace Observer object to torch_npu profiler (based on acl_prof). #44622
已合并
共 2 个文件变更+158-0
| @@ -0,0 +1,151 @@ | |||
| 1 | +# Owner(s): ["oncall: profiler"] | ||
| 2 | + | ||
| 3 | +import json | ||
| 4 | +import os | ||
| 5 | +import tempfile | ||
| 6 | +import glob | ||
| 7 | +import gzip | ||
| 8 | +from typing import Any | ||
| 9 | +import torch_npu | ||
| 10 | +import numpy as np | ||
| 11 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 12 | +from torch.autograd import ( | ||
| 13 | + _record_function_with_args_enter, | ||
| 14 | + _record_function_with_args_exit, | ||
| 15 | +) | ||
| 16 | +import torch | ||
| 17 | +import torch.nn as nn | ||
| 18 | +from torch.profiler import ExecutionTraceObserver | ||
| 19 | +from torch.autograd.profiler import record_function | ||
| 20 | +from unittest.mock import patch | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +worker_id = 1 | ||
| 24 | +Json = dict[str, Any] | ||
| 25 | + | ||
| 26 | +class TestNpuExecutionTrace(TestCase): | ||
| 27 | + """ | ||
| 28 | + NestedTensor Layout on Ascend is not supported yet. | ||
| 29 | + """ | ||
| 30 | + def trace_root(self, out_file_name) -> Json: | ||
| 31 | + nodes = [] | ||
| 32 | + with ( | ||
| 33 | + gzip.open(out_file_name) | ||
| 34 | + if out_file_name.endswith(".gz") | ||
| 35 | + else open(out_file_name) | ||
| 36 | + ) as f: | ||
| 37 | + et_graph = json.load(f) | ||
| 38 | + if "nodes" not in et_graph: | ||
| 39 | + raise AssertionError(f"Missing 'nodes' in execution trace: {et_graph}") | ||
| 40 | + nodes = et_graph["nodes"] | ||
| 41 | + return nodes | ||
| 42 | + | ||
| 43 | + def workload(self): | ||
| 44 | + device = torch.device("npu:0") | ||
| 45 | + ut = torch.randn(3, 4, 5, requires_grad=True) | ||
| 46 | + with record_function("## TEST 1 ##", "1, 2, 3"): | ||
| 47 | + t1 = torch.randn(10, 10, device=device, requires_grad=True) | ||
| 48 | + t2 = torch.randn(10, 10, device=device, requires_grad=True) | ||
| 49 | + t3 = t1 + t2 | ||
| 50 | + t3.backward(t3) | ||
| 51 | + gelu = nn.GELU() | ||
| 52 | + tm = torch.randn(2) | ||
| 53 | + _ = gelu(tm) | ||
| 54 | + t3 = t3.cpu() | ||
| 55 | + rec_fun_handler = _record_function_with_args_enter( | ||
| 56 | + "## TEST 2 ##", | ||
| 57 | + 1, | ||
| 58 | + False, | ||
| 59 | + 2.5, | ||
| 60 | + [ut, ut], | ||
| 61 | + (ut, ut), | ||
| 62 | + "hi", | ||
| 63 | + ut, | ||
| 64 | + float("inf"), float("-inf"), float("nan") | ||
| 65 | + ) | ||
| 66 | + _record_function_with_args_exit(rec_fun_handler) | ||
| 67 | + | ||
| 68 | + | ||
| 69 | + def worker_name(self): | ||
| 70 | + global worker_id | ||
| 71 | + worker_name = f"npu_profiler_test{worker_id}" | ||
| 72 | + worker_id += 1 | ||
| 73 | + return worker_name | ||
| 74 | + | ||
| 75 | + | ||
| 76 | + os.environ, | ||
| 77 | + {"ENABLE_PYTORCH_EXECUTION_TRACE_SAVE_INTEGRAL_TENSOR_RANGE": "1"}, | ||
| 78 | + ) | ||
| 79 | + def test_npu_execution_trace_record_integral_tensor_range(self): | ||
| 80 | + device = torch.device("npu:0") | ||
| 81 | + x = torch.tensor([[1, 2], [3, 4]], device=device) | ||
| 82 | + y = torch.tensor([[0, 0], [1, 0]], device=device) | ||
| 83 | + with tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False) as trace_file: | ||
| 84 | + filename = trace_file.name | ||
| 85 | + et = ExecutionTraceObserver() | ||
| 86 | + et.register_callback(filename) | ||
| 87 | + with torch_npu.profiler.profile( | ||
| 88 | + activities=[torch_npu.profiler.ProfilerActivity.CPU, | ||
| 89 | + torch_npu.profiler.ProfilerActivity.NPU], | ||
| 90 | + schedule=torch_npu.profiler.schedule( | ||
| 91 | + skip_first=0, wait=0, warmup=0, active=1, repeat=1 | ||
| 92 | + ), | ||
| 93 | + record_shapes=True, | ||
| 94 | + execution_trace_observer=et | ||
| 95 | + ) as prof: | ||
| 96 | + torch.gather(x, 1, y) | ||
| 97 | + prof.step() | ||
| 98 | + et.unregister_callback() | ||
| 99 | + nodes = self.trace_root(filename) | ||
| 100 | + os.remove(filename) | ||
| 101 | + for n in nodes: | ||
| 102 | + if "name" not in n: | ||
| 103 | + raise AssertionError(f"Expected node to have 'name': {n}") | ||
| 104 | + target_range = '{"0":[1,4],"1":[0,1]}' | ||
| 105 | + if "aten::gather" in n["name"]: | ||
| 106 | + for attr in n["attrs"]: | ||
| 107 | + if attr["name"] == "tensor_range" and attr["value"] != target_range: | ||
| 108 | + raise AssertionError(f"Expected tensor_range value to match {target_range}") | ||
| 109 | + | ||
| 110 | + def test_npu_execution_trace_record_integral_tensor_data(self): | ||
| 111 | + with tempfile.TemporaryDirectory() as temp_dir: | ||
| 112 | + fp_name = os.path.join(temp_dir, "test.et.json") | ||
| 113 | + | ||
| 114 | + os.environ["ENABLE_PYTORCH_EXECUTION_TRACE_SAVE_INTEGRAL_TENSOR_DATA"] = ( | ||
| 115 | + "aten::gather" | ||
| 116 | + ) | ||
| 117 | + et = ExecutionTraceObserver() | ||
| 118 | + et.register_callback(fp_name) | ||
| 119 | + et.set_extra_resource_collection(True) | ||
| 120 | + | ||
| 121 | + device = torch.device("npu:0") | ||
| 122 | + t1 = torch.tensor([[1, 2], [3, 4]], device=device) | ||
| 123 | + t2 = torch.tensor([[0, 0], [1, 0]], device=device) | ||
| 124 | + with torch_npu.profiler.profile( | ||
| 125 | + activities=[torch_npu.profiler.ProfilerActivity.CPU, | ||
| 126 | + torch_npu.profiler.ProfilerActivity.NPU], | ||
| 127 | + schedule=torch_npu.profiler.schedule( | ||
| 128 | + skip_first=0, wait=0, warmup=0, active=1, repeat=1 | ||
| 129 | + ), | ||
 | |||
| 130 | + record_shapes=True, | ||
| 131 | + execution_trace_observer=et, | ||
| 132 | + ) as p: | ||
| 133 | + torch.gather(t1, 1, t2) | ||
| 134 | + p.step() | ||
| 135 | + et.unregister_callback() | ||
| 136 | + | ||
| 137 | + resourceDir = fp_name.replace(".json", "_resources") | ||
| 138 | + dat_files = sorted(glob.glob(os.path.join(resourceDir, "*.dat"))) | ||
| 139 | + if len(dat_files) < 2: | ||
| 140 | + raise AssertionError(f"Expected at least 2 .dat files in {resourceDir}, found {len(dat_files)}") | ||
| 141 | + | ||
| 142 | + dumped_t1 = np.fromfile(dat_files[0], dtype=np.int64) | ||
| 143 | + dumped_t2 = np.fromfile(dat_files[1], dtype=np.int64) | ||
| 144 | + | ||
| 145 | + if not (dumped_t1 == np.array([1, 2, 3, 4])).all(): | ||
| 146 | + raise AssertionError("Expected t1 contents to match [1, 2, 3, 4]") | ||
| 147 | + if not (dumped_t2 == np.array([0, 0, 1, 0])).all(): | ||
| 148 | + raise AssertionError("Expected t2 contents to match [0, 0, 1, 0]") | ||
| 149 | + | ||
| 150 | +if __name__ == "__main__": | ||
| 151 | + run_tests() | ||
| @@ -7,6 +7,7 @@ from typing import Any, Optional, Union | |||
| 7 | 7 | ||
| 8 | import torch.autograd.profiler as prof | 8 | import torch.autograd.profiler as prof |
| 9 | import torch_npu.npu | 9 | import torch_npu.npu |
| 10 | +from torch.profiler.profiler import _ITraceObserver | ||
| 10 | from torch_npu.npu import current_stream, mstx | 11 | from torch_npu.npu import current_stream, mstx |
| 11 | from torch_npu._C._profiler import ( | 12 | from torch_npu._C._profiler import ( |
| 12 | _disable_profiler_in_child_thread, | 13 | _disable_profiler_in_child_thread, |
| @@ -211,6 +212,7 @@ class profile(_KinetoProfile): | |||
| 211 | with_flops: bool = False, | 212 | with_flops: bool = False, |
| 212 | with_modules: bool = False, | 213 | with_modules: bool = False, |
| 213 | experimental_config: Optional[_ExperimentalConfig] = None, | 214 | experimental_config: Optional[_ExperimentalConfig] = None, |
| 215 | + execution_trace_observer: _ITraceObserver | None = None, | ||
| 214 | custom_trace_id_callback: Optional[Callable[[], str]] = None, | 216 | custom_trace_id_callback: Optional[Callable[[], str]] = None, |
| 215 | # deprecated: | 217 | # deprecated: |
| 216 | use_cuda: Optional[bool] = None, | 218 | use_cuda: Optional[bool] = None, |
| @@ -249,6 +251,7 @@ class profile(_KinetoProfile): | |||
| 249 | metadata=self.metadata, | 251 | metadata=self.metadata, |
| 250 | custom_trace_id_callback=custom_trace_id_callback, | 252 | custom_trace_id_callback=custom_trace_id_callback, |
| 251 | ) | 253 | ) |
| 254 | + self.execution_trace_observer = execution_trace_observer | ||
| 252 | self.on_trace_ready = on_trace_ready | 255 | self.on_trace_ready = on_trace_ready |
| 253 | self.step_num = 0 | 256 | self.step_num = 0 |
| 254 | self.current_action = self.schedule(self.step_num) | 257 | self.current_action = self.schedule(self.step_num) |
| @@ -303,6 +306,8 @@ class profile(_KinetoProfile): | |||
| 303 | self.step_rec_fn = prof.record_function(step_name) | 306 | self.step_rec_fn = prof.record_function(step_name) |
| 304 | self.step_rec_fn.__enter__() | 307 | self.step_rec_fn.__enter__() |
| 305 | self._start_step_mstx_range(step_name) | 308 | self._start_step_mstx_range(step_name) |
| 309 | + if self.execution_trace_observer: | ||
| 310 | + self.execution_trace_observer.start() | ||
| 306 | 311 | ||
| 307 | 312 | ||
| 308 | def stop(self): | 313 | def stop(self): |
| @@ -311,6 +316,8 @@ class profile(_KinetoProfile): | |||
| 311 | self._end_step_mstx_range() | 316 | self._end_step_mstx_range() |
| 312 | self.action_controller.transit_action(self.current_action, None) | 317 | self.action_controller.transit_action(self.current_action, None) |
| 313 | self.stopped = True | 318 | self.stopped = True |
| 319 | + if self.execution_trace_observer: | ||
| 320 | + self.execution_trace_observer.stop() | ||
| 314 | 321 | ||
| 315 | 322 | ||
| 316 | def step(self): | 323 | def step(self): |
🟡 Medium Priority
建议：将第 127 行的
schedule=torch.profiler.schedule(...)改为schedule=torch_npu.profiler.schedule(...)，使用与第一个测试及 torch_npu profiler 内部约定一致的调度器。