已合并
【pytorch_v2.10.0】 add dynamic step #37366
hewenbo创建于 6月1日
【pytorch_v2.10.0】 add dynamic step #37366
已合并
共 3 个文件变更+57-5
| @@ -798,9 +798,18 @@ class TestDynamicProfiler(TestCase): | |||
| 798 | self.assertIsNone(result) | 798 | self.assertIsNone(result) |
| 799 | 799 | ||
| 800 | def test_start_while_profiler_active(self): | 800 | def test_start_while_profiler_active(self): |
| 801 | - dp.start() | 801 | + cfg_json = copy.deepcopy(self.json_sample) |
| 802 | - dp.start() | 802 | + cfg_json["prof_dir"] = self.active_rank_prof_dir |
| 803 | + with os.fdopen(os.open(self.cfg_path, self.flags, self.mode), "w") as f: | ||
| 804 | + json.dump(cfg_json, f, indent=4) | ||
| 805 | + | ||
| 806 | + dp.start(self.cfg_path) | ||
| 807 | + dp.start(self.cfg_path) | ||
| 803 | self.assertIsNotNone(_DynamicProfile().prof) | 808 | self.assertIsNotNone(_DynamicProfile().prof) |
| 809 | + _DynamicProfile().prof.stop() | ||
| 810 | + _DynamicProfile().prof = None | ||
| 811 | + if os.path.exists(self.active_rank_prof_dir): | ||
| 812 | + PathManager.remove_path_safety(self.active_rank_prof_dir) | ||
| 804 | 813 | ||
| 805 | def test_init_repeated_warning(self): | 814 | def test_init_repeated_warning(self): |
| 806 | dp.init(self.results_path) | 815 | dp.init(self.results_path) |
| @@ -815,6 +824,28 @@ class TestDynamicProfiler(TestCase): | |||
| 815 | dynamic_prof.step() | 824 | dynamic_prof.step() |
| 816 | self.assertIsNotNone(dynamic_prof._step_time) | 825 | self.assertIsNotNone(dynamic_prof._step_time) |
| 817 | 826 | ||
| 827 | + def test_set_state(self): | ||
| 828 | + dynamic_prof = _DynamicProfile() | ||
| 829 | + dynamic_prof.cur_step = 7 | ||
| 830 | + | ||
| 831 | + dynamic_prof.cur_step = 0 | ||
| 832 | + dynamic_prof.set_state({"cur_step": 7}) | ||
| 833 | + self.assertEqual(dynamic_prof.cur_step, 7) | ||
| 834 | + | ||
| 835 | + dynamic_prof.set_state({"cur_step": -1}) | ||
| 836 | + self.assertEqual(dynamic_prof.cur_step, 7) | ||
| 837 | + | ||
| 838 | + def test_set_state_continue_step(self): | ||
| 839 | + dynamic_prof = _DynamicProfile() | ||
| 840 | + dynamic_prof.cur_step = 0 | ||
| 841 | + dynamic_prof._dynamic_monitor = MagicMock() | ||
| 842 | + dynamic_prof._dynamic_monitor.shm_to_prof_conf_context.return_value = None | ||
| 843 | + dynamic_prof._step_mstx_range_id = 0 | ||
| 844 | + dynamic_prof.set_state({"cur_step": 10}) | ||
| 845 | + with patch("torch_npu.profiler.dynamic_profile.mstx.range_start", return_value=1): | ||
| 846 | + dynamic_prof.step() | ||
| 847 | + self.assertEqual(dynamic_prof.cur_step, 11) | ||
| 848 | + | ||
| 818 | 849 | ||
| 819 | if __name__ == "__main__": | 850 | if __name__ == "__main__": |
| 820 | run_tests() | 851 | run_tests() |
| @@ -1820,6 +1820,9 @@ | |||
| 1820 | "torch_npu.profiler.dynamic_profile.start": { | 1820 | "torch_npu.profiler.dynamic_profile.start": { |
| 1821 | "signature": "(config_path: str = None)" | 1821 | "signature": "(config_path: str = None)" |
| 1822 | }, | 1822 | }, |
| 1823 | + "torch_npu.profiler.dynamic_profile.set_state": { | ||
| 1824 | + "signature": "(state_step: dict)" | ||
| 1825 | + }, | ||
| 1823 | "torch_npu.profiler.experimental_config.AiCMetrics": { | 1826 | "torch_npu.profiler.experimental_config.AiCMetrics": { |
| 1824 | "signature": "()" | 1827 | "signature": "()" |
| 1825 | }, | 1828 | }, |
| @@ -2904,4 +2907,4 @@ | |||
| 2904 | "signature": "(std::vector<std::string>& op_type, std::vector<at::Tensor>& tensors, std::vector<int64_t> remote_rank_list) -> c10::intrusive_ptr<c10d::Work>", | 2907 | "signature": "(std::vector<std::string>& op_type, std::vector<at::Tensor>& tensors, std::vector<int64_t> remote_rank_list) -> c10::intrusive_ptr<c10d::Work>", |
| 2905 | "file": "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp" | 2908 | "file": "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp" |
| 2906 | } | 2909 | } |
| 2907 | -} | 2910 | +} |
| @@ -21,7 +21,8 @@ from ._dynamic_profiler._dynamic_monitor_proxy import PyDynamicMonitorProxySingl | |||
| 21 | __all__ = [ | 21 | __all__ = [ |
| 22 | 'init', | 22 | 'init', |
| 23 | 'step', | 23 | 'step', |
| 24 | - 'start' | 24 | + 'start', |
| 25 | + 'set_state' | ||
| 25 | ] | 26 | ] |
| 26 | 27 | ||
| 27 | 28 | ||
| @@ -71,6 +72,18 @@ class _DynamicProfile: | |||
| 71 | prof_cfg_ctx = self._dynamic_monitor.shm_to_prof_conf_context() | 72 | prof_cfg_ctx = self._dynamic_monitor.shm_to_prof_conf_context() |
| 72 | return prof_cfg_ctx | 73 | return prof_cfg_ctx |
| 73 | 74 | ||
| 75 | + def set_state(self, state_step: dict): | ||
| 76 | + if not isinstance(state_step, dict): | ||
| 77 | + DynamicProfilerUtils.stdout_log("Dynamic profiler state_step must be dict.", | ||
| 78 | + DynamicProfilerUtils.LoggerLevelEnum.ERROR) | ||
| 79 | + return | ||
| 80 | + cur_step = state_step.get("cur_step", 0) | ||
| 81 | + if not isinstance(cur_step, int) or cur_step < 0: | ||
| 82 | + DynamicProfilerUtils.stdout_log("Dynamic profiler cur_step must be a non-negative integer.", | ||
| 83 | + DynamicProfilerUtils.LoggerLevelEnum.ERROR) | ||
| 84 | + return | ||
| 85 | + self.cur_step = cur_step | ||
| 86 | + | ||
| 74 | def step(self): | 87 | def step(self): |
| 75 | self.cur_step += 1 | 88 | self.cur_step += 1 |
| 76 | cfg_ctx = self._dynamic_profiler_valid() | 89 | cfg_ctx = self._dynamic_profiler_valid() |
| @@ -78,7 +91,7 @@ class _DynamicProfile: | |||
| 78 | self.cfg_ctx = cfg_ctx | 91 | self.cfg_ctx = cfg_ctx |
| 79 | if self.cur_step == self.RECORD_TIME_STEP: | 92 | if self.cur_step == self.RECORD_TIME_STEP: |
| 80 | self._step_record_time = time.time() | 93 | self._step_record_time = time.time() |
| 81 | - elif self.cur_step - self.RECORD_TIME_STEP == 1: | 94 | + elif self.cur_step - self.RECORD_TIME_STEP == 1 and self._step_record_time is not None: |
| 82 | self._step_time = min(self._max_poll_interval, | 95 | self._step_time = min(self._max_poll_interval, |
| 83 | max(self._min_poll_interval, int(time.time() - self._step_record_time))) | 96 | max(self._min_poll_interval, int(time.time() - self._step_record_time))) |
| 84 | self._dynamic_monitor.modify_step_time(self._step_time) | 97 | self._dynamic_monitor.modify_step_time(self._step_time) |
| @@ -201,3 +214,8 @@ def step(): | |||
| 201 | 214 | ||
| 202 | def start(config_path: str = None): | 215 | def start(config_path: str = None): |
| 203 | _DynamicProfile().start(config_path) | 216 | _DynamicProfile().start(config_path) |
| 217 | + | ||
| 218 | + | ||
| 219 | + | ||
| 220 | +def set_state(state_step: dict): | ||
| 221 | + _DynamicProfile().set_state(state_step) | ||