已合并
【优化】增加debug调试日志，部分info/warning级别日志切换为rank_0输出 #873
feQQ创建于 7月6日
【优化】增加debug调试日志，部分info/warning级别日志切换为rank_0输出 #873
已合并
共 9 个文件变更+130-63
| @@ -37,6 +37,13 @@ class CommonConfig: | |||
| 37 | self.risk_level = json_config.get("risk_level", Const.RISK_LEVEL_FOCUS) | 37 | self.risk_level = json_config.get("risk_level", Const.RISK_LEVEL_FOCUS) |
| 38 | self.custom_op_namespaces = json_config.get("custom_op_namespaces") | 38 | self.custom_op_namespaces = json_config.get("custom_op_namespaces") |
| 39 | self._check_config() | 39 | self._check_config() |
| 40 | + logger.debug( | ||
| 41 | + f"CommonConfig: task={self.task}, dump_path={self.dump_path}, " | ||
| 42 | + f"rank={self.rank}, step={self.step}, level={self.level}, " | ||
| 43 | + f"dump_enable={self.dump_enable}, extra_info={self.extra_info}, " | ||
| 44 | + f"async_dump={self.async_dump}, precision={self.precision}, " | ||
| 45 | + f"risk_level={self.risk_level}, custom_op_namespaces={self.custom_op_namespaces}" | ||
| 46 | + ) | ||
| 40 | 47 | ||
| 41 | def _check_config(self): | 48 | def _check_config(self): |
| 42 | if self.task and self.task not in Const.TASK_LIST: | 49 | if self.task and self.task not in Const.TASK_LIST: |
| @@ -100,6 +107,11 @@ class BaseConfig: | |||
| 100 | self.summary_mode = json_config.get("summary_mode") | 107 | self.summary_mode = json_config.get("summary_mode") |
| 101 | self.diff_nums = json_config.get("diff_nums") | 108 | self.diff_nums = json_config.get("diff_nums") |
| 102 | self.is_regex_valid = True | 109 | self.is_regex_valid = True |
| 110 | + logger.debug( | ||
| 111 | + f"BaseConfig: scope={self.scope}, list={self.list}, " | ||
| 112 | + f"data_mode={self.data_mode}, summary_mode={self.summary_mode}, " | ||
| 113 | + f"diff_nums={self.diff_nums}" | ||
| 114 | + ) | ||
| 103 | 115 | ||
| 104 | @staticmethod | 116 | @staticmethod |
| 105 | def _check_str_list_config(config_item, config_name): | 117 | def _check_str_list_config(config_item, config_name): |
| @@ -90,6 +90,7 @@ class DataCollector: | |||
| 90 | return True | 90 | return True |
| 91 | if self.config.framework == Const.PT_FRAMEWORK: | 91 | if self.config.framework == Const.PT_FRAMEWORK: |
| 92 | from msprobe.pytorch.dump.api_dump.api_risk_level import get_api_risk_level | 92 | from msprobe.pytorch.dump.api_dump.api_risk_level import get_api_risk_level |
| 93 | + | ||
| 93 | api_name = self._api_name_pattern.sub('', name).split(Const.SEP, 1) | 94 | api_name = self._api_name_pattern.sub('', name).split(Const.SEP, 1) |
| 94 | prefix = api_name[0] | 95 | prefix = api_name[0] |
| 95 | if prefix in Const.MODULE_PREFIX: | 96 | if prefix in Const.MODULE_PREFIX: |
| @@ -136,6 +137,7 @@ class DataCollector: | |||
| 136 | if not self.should_collect(name, pid): | 137 | if not self.should_collect(name, pid): |
| 137 | return | 138 | return |
| 138 | 139 | ||
| 140 | + logger.debug(f"forward_input_data_collect for {name}") | ||
| 139 | data_info = {} | 141 | data_info = {} |
| 140 | if self.config.task != Const.STRUCTURE: | 142 | if self.config.task != Const.STRUCTURE: |
| 141 | data_info = self.data_processor.analyze_forward_input(name, module, module_input_output) | 143 | data_info = self.data_processor.analyze_forward_input(name, module, module_input_output) |
| @@ -145,12 +147,11 @@ class DataCollector: | |||
| 145 | self.handle_data(name, data_info, flush=self.data_processor.is_terminated) | 147 | self.handle_data(name, data_info, flush=self.data_processor.is_terminated) |
| 146 | 148 | ||
| 147 | except Exception as e: | 149 | except Exception as e: |
| 148 | - # 取异常类名作为“类型”做去重 | ||
| 149 | error_type = type(e).__name__ | 150 | error_type = type(e).__name__ |
| 150 | tb = traceback.format_exc() | 151 | tb = traceback.format_exc() |
| 152 | + logger.debug(f"forward_input_data_collect failed for {name}: {e}") | ||
| 151 | self.data_writer.write_error_log( | 153 | self.data_writer.write_error_log( |
| 152 | - f"[ERROR] forward_input_data_collect failed: name={name}, pid={pid}\n{tb}", | 154 | + f"[ERROR] forward_input_data_collect failed: name={name}, pid={pid}\n{tb}", error_type=error_type |
| 153 | - error_type=error_type | ||
| 154 | ) | 155 | ) |
| 155 | 156 | ||
| 156 | def forward_output_data_collect(self, name, module, pid, module_input_output): | 157 | def forward_output_data_collect(self, name, module, pid, module_input_output): |
| @@ -158,6 +159,7 @@ class DataCollector: | |||
| 158 | if not self.should_collect(name, pid): | 159 | if not self.should_collect(name, pid): |
| 159 | return | 160 | return |
| 160 | 161 | ||
| 162 | + logger.debug(f"forward_output_data_collect for {name}") | ||
| 161 | data_info = {} | 163 | data_info = {} |
| 162 | if self.config.task != Const.STRUCTURE: | 164 | if self.config.task != Const.STRUCTURE: |
| 163 | data_info = self.data_processor.analyze_forward_output(name, module, module_input_output) | 165 | data_info = self.data_processor.analyze_forward_output(name, module, module_input_output) |
| @@ -189,6 +191,7 @@ class DataCollector: | |||
| 189 | self.update_construct(name) | 191 | self.update_construct(name) |
| 190 | if not self.should_collect(name, pid): | 192 | if not self.should_collect(name, pid): |
| 191 | return | 193 | return |
| 194 | + logger.debug(f"backward_data_collect for {name}") | ||
| 192 | data_info = {} | 195 | data_info = {} |
| 193 | if self.config.task != Const.STRUCTURE: | 196 | if self.config.task != Const.STRUCTURE: |
| 194 | data_info = self.data_processor.analyze_backward(name, module, module_input_output) | 197 | data_info = self.data_processor.analyze_backward(name, module, module_input_output) |
| @@ -224,13 +227,16 @@ class DataCollector: | |||
| 224 | if self.optimizer_status in [Const.OPTIMIZER, Const.CLIP_GRAD]: | 227 | if self.optimizer_status in [Const.OPTIMIZER, Const.CLIP_GRAD]: |
| 225 | if self.optimizer_status_first_start[self.optimizer_status]: | 228 | if self.optimizer_status_first_start[self.optimizer_status]: |
| 226 | self.data_writer.update_construct( | 229 | self.data_writer.update_construct( |
| 227 | - {self.optimizer_status: None if not is_megatron() else [None, get_micro_step()]}) | 230 | + {self.optimizer_status: None if not is_megatron() else [None, get_micro_step()]} |
| 231 | + ) | ||
| 228 | self.optimizer_status_first_start[self.optimizer_status] = False | 232 | self.optimizer_status_first_start[self.optimizer_status] = False |
| 229 | self.data_writer.update_construct( | 233 | self.data_writer.update_construct( |
| 230 | - {name: self.optimizer_status if not is_megatron() else [self.optimizer_status, get_micro_step()]}) | 234 | + {name: self.optimizer_status if not is_megatron() else [self.optimizer_status, get_micro_step()]} |
| 235 | + ) | ||
| 231 | else: | 236 | else: |
| 232 | - if self.config.level == Const.LEVEL_MIX and \ | 237 | + if self.config.level == Const.LEVEL_MIX and not ( |
| 233 | - not (name.startswith(Const.MODULE) or name.startswith(Const.CELL)): | 238 | + name.startswith(Const.MODULE) or name.startswith(Const.CELL) |
| 239 | + ): | ||
| 234 | self.data_writer.update_construct( | 240 | self.data_writer.update_construct( |
| 235 | {name: self.module_processor.api_parent_node.get(threading.get_ident())} | 241 | {name: self.module_processor.api_parent_node.get(threading.get_ident())} |
| 236 | ) | 242 | ) |
| @@ -242,6 +248,7 @@ class DataCollector: | |||
| 242 | 248 | ||
| 243 | def handle_data(self, name, data_info, flush=False): | 249 | def handle_data(self, name, data_info, flush=False): |
| 244 | if data_info: | 250 | if data_info: |
| 251 | + logger.debug(f"handle_data: saving data for {name} to cache") | ||
| 245 | self.update_data(name, data_info) | 252 | self.update_data(name, data_info) |
| 246 | if self.config.async_dump: | 253 | if self.config.async_dump: |
| 247 | return | 254 | return |
| @@ -288,5 +295,6 @@ class DataCollector: | |||
| 288 | self.data_writer.update_debug({grad_name_with_count_category: all_none_data_info}) | 295 | self.data_writer.update_debug({grad_name_with_count_category: all_none_data_info}) |
| 289 | 296 | ||
| 290 | # register tensor backward hook | 297 | # register tensor backward hook |
| 291 | - self.data_processor.analyze_debug_backward(variable, grad_name_with_count_category, | 298 | + self.data_processor.analyze_debug_backward( |
| 292 | - self.data_writer.cache_debug['data']) | 299 | + variable, grad_name_with_count_category, self.data_writer.cache_debug['data'] |
| 300 | + ) | ||
| @@ -313,7 +313,7 @@ class DataWriter: | |||
| 313 | self.cache_debug['data'].update(new_data) | 313 | self.cache_debug['data'].update(new_data) |
| 314 | 314 | ||
| 315 | def write_data_json(self, file_path): | 315 | def write_data_json(self, file_path): |
| 316 | - logger.info(f"dump.json is at {os.path.dirname(os.path.dirname(file_path))}. ") | 316 | + logger.info_on_rank_0(f"dump.json is at {os.path.dirname(os.path.dirname(file_path))}. ") |
| 317 | save_json(file_path, self.cache_data, indent=1) | 317 | save_json(file_path, self.cache_data, indent=1) |
| 318 | 318 | ||
| 319 | def write_stack_info_json(self, file_path): | 319 | def write_stack_info_json(self, file_path): |
| @@ -405,11 +405,15 @@ class DataWriter: | |||
| 405 | return | 405 | return |
| 406 | 406 | ||
| 407 | if self.cache_data: | 407 | if self.cache_data: |
| 408 | + data_count = len(self.cache_data.get(Const.DATA, {})) | ||
| 409 | + logger.debug(f"write_json: dump.json written with {data_count} entries") | ||
| 408 | self.write_data_json(self.dump_file_path) | 410 | self.write_data_json(self.dump_file_path) |
| 409 | if self._is_extra_info_enabled() and self.cache_stack: | 411 | if self._is_extra_info_enabled() and self.cache_stack: |
| 410 | self.write_stack_info_json(self.stack_file_path) | 412 | self.write_stack_info_json(self.stack_file_path) |
| 411 | if self._is_extra_info_enabled() and self.cache_construct: | 413 | if self._is_extra_info_enabled() and self.cache_construct: |
| 412 | self.write_construct_info_json(self.construct_file_path) | 414 | self.write_construct_info_json(self.construct_file_path) |
| 413 | if self.cache_debug: | 415 | if self.cache_debug: |
| 416 | + debug_count = len(self.cache_debug.get(Const.DATA, {})) | ||
| 417 | + logger.debug(f"write_json: debug.json written with {debug_count} entries") | ||
| 414 | self.write_debug_info_json(self.debug_file_path) | 418 | self.write_debug_info_json(self.debug_file_path) |
| 415 | self.data_updated = False | 419 | self.data_updated = False |
| @@ -213,6 +213,8 @@ class BaseHookManager(ABC): | |||
| 213 | if not self._should_execute_hook(hook_type, tid): | 213 | if not self._should_execute_hook(hook_type, tid): |
| 214 | return None | 214 | return None |
| 215 | with ThreadSafe(): | 215 | with ThreadSafe(): |
| 216 | + if self.logger: | ||
| 217 | + self.logger.debug(f"forward_pre_hook executed for {api_name}, type={hook_type}") | ||
| 216 | self._on_forward_pre_hook() | 218 | self._on_forward_pre_hook() |
| 217 | self._maybe_update_dump_dir() | 219 | self._maybe_update_dump_dir() |
| 218 | original_state = self.ensure_gc_enabled() | 220 | original_state = self.ensure_gc_enabled() |
| @@ -249,6 +251,8 @@ class BaseHookManager(ABC): | |||
| 249 | return None | 251 | return None |
| 250 | 252 | ||
| 251 | with ThreadSafe(): | 253 | with ThreadSafe(): |
| 254 | + if self.logger: | ||
| 255 | + self.logger.debug(f"forward_hook executed for {api_name}, type={hook_type}") | ||
| 252 | self._maybe_update_dump_dir() | 256 | self._maybe_update_dump_dir() |
| 253 | original_state = self.ensure_gc_enabled() | 257 | original_state = self.ensure_gc_enabled() |
| 254 | if hook_type == Const.API: | 258 | if hook_type == Const.API: |
| @@ -298,6 +302,8 @@ class BaseHookManager(ABC): | |||
| 298 | return | 302 | return |
| 299 | 303 | ||
| 300 | with ThreadSafe(): | 304 | with ThreadSafe(): |
| 305 | + if self.logger: | ||
| 306 | + self.logger.debug(f"backward_hook executed for {full_name}, type={hook_type}") | ||
| 301 | self._maybe_update_dump_dir() | 307 | self._maybe_update_dump_dir() |
| 302 | original_state = self.ensure_gc_enabled() | 308 | original_state = self.ensure_gc_enabled() |
| 303 | BaseHookManager.inner_switch[tid] = True | 309 | BaseHookManager.inner_switch[tid] = True |
| @@ -95,8 +95,8 @@ class BaseService(ABC): | |||
| 95 | def _need_tensor_data(self): | 95 | def _need_tensor_data(self): |
| 96 | """Whether tensor data collection is required.""" | 96 | """Whether tensor data collection is required.""" |
| 97 | return bool( | 97 | return bool( |
| 98 | - self.config.task in self.data_collector.tasks_need_tensor_data or | 98 | + self.config.task in self.data_collector.tasks_need_tensor_data |
| 99 | - (self.config.task == Const.STATISTICS and self.config.tensor_list) | 99 | + or (self.config.task == Const.STATISTICS and self.config.tensor_list) |
| 100 | ) | 100 | ) |
| 101 | 101 | ||
| 102 | @property | 102 | @property |
| @@ -155,14 +155,14 @@ class BaseService(ABC): | |||
| 155 | elif self._is_no_dump_rank: | 155 | elif self._is_no_dump_rank: |
| 156 | Runtime.is_running = False | 156 | Runtime.is_running = False |
| 157 | return | 157 | return |
| 158 | - self.logger.info(f"{Const.TOOL_NAME}: debugger.start() is set successfully") | 158 | + self.logger.info_on_rank_0(f"{Const.TOOL_NAME}: debugger.start() is set successfully") |
| 159 | if token_range is None: | 159 | if token_range is None: |
| 160 | self.primitive_switch = True | 160 | self.primitive_switch = True |
| 161 | self._change_jit_switch(True) | 161 | self._change_jit_switch(True) |
| 162 | - self.logger.info(f"Dump switch is turned on at step {self.current_iter}. ") | 162 | + self.logger.info_on_rank_0(f"Dump switch is turned on at step {self.current_iter}. ") |
| 163 | 163 | ||
| 164 | self.create_dirs() | 164 | self.create_dirs() |
| 165 | - self.logger.info(f"Dump data will be saved in {self.dump_iter_dir}.") | 165 | + self.logger.info_on_rank_0(f"Dump data will be saved in {self.dump_iter_dir}.") |
| 166 | 166 | ||
| 167 | def stop(self): | 167 | def stop(self): |
| 168 | """Common stop flow.""" | 168 | """Common stop flow.""" |
| @@ -173,8 +173,10 @@ class BaseService(ABC): | |||
| 173 | return | 173 | return |
| 174 | if self._is_no_dump_step or self._is_no_dump_rank: | 174 | if self._is_no_dump_step or self._is_no_dump_rank: |
| 175 | return | 175 | return |
| 176 | - self.logger.info(f"{Const.TOOL_NAME}: debugger.stop() is set successfully. " | 176 | + self.logger.info_on_rank_0( |
| 177 | - "Please set debugger.start() to turn on the dump switch again. ") | 177 | + f"{Const.TOOL_NAME}: debugger.stop() is set successfully. " |
| 178 | + "Please set debugger.start() to turn on the dump switch again. " | ||
| 179 | + ) | ||
| 178 | Runtime.is_running = False | 180 | Runtime.is_running = False |
| 179 | self.primitive_switch = False | 181 | self.primitive_switch = False |
| 180 | self._change_jit_switch(False) | 182 | self._change_jit_switch(False) |
| @@ -234,8 +236,9 @@ class BaseService(ABC): | |||
| 234 | 236 | ||
| 235 | def register_custom_api(self, module, api_name, api_prefix): | 237 | def register_custom_api(self, module, api_name, api_prefix): |
| 236 | self.ori_customer_func[str(module) + Const.SEP + api_name] = getattr(module, api_name) | 238 | self.ori_customer_func[str(module) + Const.SEP + api_name] = getattr(module, api_name) |
| 237 | - ApiRegistry.register_custom_api(module, api_name, api_prefix, | 239 | + ApiRegistry.register_custom_api( |
| 238 | - functools.partial(self.build_hook, Const.API), self.api_template) | 240 | + module, api_name, api_prefix, functools.partial(self.build_hook, Const.API), self.api_template |
| 241 | + ) | ||
| 239 | 242 | ||
| 240 | def restore_custom_api(self, module, api): | 243 | def restore_custom_api(self, module, api): |
| 241 | ori_func = self.ori_customer_func.get(str(module) + Const.SEP + api) | 244 | ori_func = self.ori_customer_func.get(str(module) + Const.SEP + api) |
| @@ -285,8 +288,12 @@ class BaseService(ABC): | |||
| 285 | def _need_stop_service(self): | 288 | def _need_stop_service(self): |
| 286 | if self.should_stop_service: | 289 | if self.should_stop_service: |
| 287 | return True | 290 | return True |
| 288 | - end_service = self.config.step and self.current_iter > max(self.config.step) or \ | 291 | + end_service = ( |
| 289 | - self.data_collector and self.data_collector.data_processor.is_terminated | 292 | + self.config.step |
| 293 | + and self.current_iter > max(self.config.step) | ||
| 294 | + or self.data_collector | ||
| 295 | + and self.data_collector.data_processor.is_terminated | ||
| 296 | + ) | ||
| 290 | if end_service: | 297 | if end_service: |
| 291 | self.primitive_switch = False | 298 | self.primitive_switch = False |
| 292 | self._change_jit_switch(False) | 299 | self._change_jit_switch(False) |
| @@ -302,7 +309,9 @@ class BaseService(ABC): | |||
| 302 | if self._is_need_api_hook: | 309 | if self._is_need_api_hook: |
| 303 | self.api_register.initialize_hook(functools.partial(self.build_hook, Const.API)) | 310 | self.api_register.initialize_hook(functools.partial(self.build_hook, Const.API)) |
| 304 | self.api_register.register_all_api() | 311 | self.api_register.register_all_api() |
| 305 | - self.logger.info(f"The api {self.config.task} hook function is successfully mounted to the model.") | 312 | + self.logger.info_on_rank_0( |
| 313 | + f"The api {self.config.task} hook function is successfully mounted to the model." | ||
| 314 | + ) | ||
| 306 | 315 | ||
| 307 | def _register_infer_count_hook(self, root_model, token_range): | 316 | def _register_infer_count_hook(self, root_model, token_range): |
| 308 | """ | 317 | """ |
| @@ -323,21 +332,24 @@ class BaseService(ABC): | |||
| 323 | Runtime.is_running = True | 332 | Runtime.is_running = True |
| 324 | self.primitive_switch = True | 333 | self.primitive_switch = True |
| 325 | self._change_jit_switch(True) | 334 | self._change_jit_switch(True) |
| 326 | - self.logger.info(f"Current token id: {self.cur_token_id}, start dump infer token.") | 335 | + self.logger.info_on_rank_0(f"Current token id: {self.cur_token_id}, start dump infer token.") |
| 327 | elif token_range[0] < self.cur_token_id <= token_range[1]: | 336 | elif token_range[0] < self.cur_token_id <= token_range[1]: |
| 328 | self.logger.debug(f"Current token id: {self.cur_token_id}.") | 337 | self.logger.debug(f"Current token id: {self.cur_token_id}.") |
| 329 | elif self.cur_token_id == token_range[1] + 1: | 338 | elif self.cur_token_id == token_range[1] + 1: |
| 330 | Runtime.is_running = False | 339 | Runtime.is_running = False |
| 331 | self.primitive_switch = False | 340 | self.primitive_switch = False |
| 332 | self._change_jit_switch(False) | 341 | self._change_jit_switch(False) |
| 333 | - self.logger.info( | 342 | + self.logger.info_on_rank_0( |
| 334 | - f"Current token id: {self.cur_token_id}, exceed token_range, early stop dump infer token.") | 343 | + f"Current token id: {self.cur_token_id}, exceed token_range, early stop dump infer token." |
| 344 | + ) | ||
| 335 | self.cur_token_id += 1 | 345 | self.cur_token_id += 1 |
| 336 | 346 | ||
| 337 | # root_model is guaranteed to be Module/Cell or [Module/Cell] | 347 | # root_model is guaranteed to be Module/Cell or [Module/Cell] |
| 338 | if root_model and isinstance(root_model, list): | 348 | if root_model and isinstance(root_model, list): |
| 339 | root_model = root_model[0] | 349 | root_model = root_model[0] |
| 340 | - self.logger.warning("Infer model can only input one to support token_range, choose the first one.") | 350 | + self.logger.warning_on_rank_0( |
| 351 | + "Infer model can only input one to support token_range, choose the first one." | ||
| 352 | + ) | ||
| 341 | 353 | ||
| 342 | root_model.register_forward_pre_hook(infer_hook) | 354 | root_model.register_forward_pre_hook(infer_hook) |
| 343 | 355 | ||
| @@ -424,4 +436,3 @@ class BaseService(ABC): | |||
| 424 | self.hook_manager.reset_status() | 436 | self.hook_manager.reset_status() |
| 425 | if self._is_l2_level: | 437 | if self._is_l2_level: |
| 426 | self.data_collector.data_processor.reset_status() | 438 | self.data_collector.data_processor.reset_status() |
| 427 | - | ||
| @@ -181,7 +181,7 @@ class PrecisionDebugger(BasePrecisionDebugger): | |||
| 181 | return False | 181 | return False |
| 182 | 182 | ||
| 183 | self._apply_reloaded_config(common_config, task_config, new_config, pending_signature) | 183 | self._apply_reloaded_config(common_config, task_config, new_config, pending_signature) |
| 184 | - logger.info("PrecisionDebugger detected config change and reloaded runtime settings.") | 184 | + logger.info_on_rank_0("PrecisionDebugger detected config change and reloaded runtime settings.") |
| 185 | return True | 185 | return True |
| 186 | 186 | ||
| 187 | # pylint: disable=attribute-defined-outside-init | 187 | # pylint: disable=attribute-defined-outside-init |
| @@ -251,7 +251,7 @@ class PrecisionDebugger(BasePrecisionDebugger): | |||
| 251 | logger.debug(f"Failed to scan custom ops namespace {ns_name}: {ex}") | 251 | logger.debug(f"Failed to scan custom ops namespace {ns_name}: {ex}") |
| 252 | 252 | ||
| 253 | if registered_count > 0: | 253 | if registered_count > 0: |
| 254 | - logger.info(f"Auto registered {registered_count} custom ops for dump.") | 254 | + logger.info_on_rank_0(f"Auto registered {registered_count} custom ops for dump.") |
| 255 | 255 | ||
| 256 | def _register_single_op(self, ns_name, op_name, ns_module): | 256 | def _register_single_op(self, ns_name, op_name, ns_module): |
| 257 | op_key = (ns_name, op_name) | 257 | op_key = (ns_name, op_name) |
| @@ -319,7 +319,7 @@ class PrecisionDebugger(BasePrecisionDebugger): | |||
| 319 | raise AttributeError(f"{module_path} does not have attribute {api_name}") | 319 | raise AttributeError(f"{module_path} does not have attribute {api_name}") |
| 320 | self.__class__.register_custom_api(module_obj, api_name, api_prefix) | 320 | self.__class__.register_custom_api(module_obj, api_name, api_prefix) |
| 321 | self._custom_api_auto_registered.add(key) | 321 | self._custom_api_auto_registered.add(key) |
| 322 | - logger.info(f"Auto-registered custom api from yaml: {module_path}.{api_name}") | 322 | + logger.info_on_rank_0(f"Auto-registered custom api from yaml: {module_path}.{api_name}") |
| 323 | except Exception as ex: | 323 | except Exception as ex: |
| 324 | pending.append(item) | 324 | pending.append(item) |
| 325 | logger.warning(f"Auto-register custom api from yaml skipped: {item}, reason: {ex}") | 325 | logger.warning(f"Auto-register custom api from yaml skipped: {item}, reason: {ex}") |
| @@ -348,7 +348,9 @@ class PrecisionDebugger(BasePrecisionDebugger): | |||
| 348 | logger.warning(f"Failed to load custom api yaml: {yaml_path}, reason: {ex}") | 348 | logger.warning(f"Failed to load custom api yaml: {yaml_path}, reason: {ex}") |
| 349 | return [] | 349 | return [] |
| 350 | if not content: | 350 | if not content: |
| 351 | - logger.info(f"Custom api yaml is empty: {yaml_path}, yaml-based custom op registration is disabled.") | 351 | + logger.info_on_rank_0( |
| 352 | + f"Custom api yaml is empty: {yaml_path}, yaml-based custom op registration is disabled." | ||
| 353 | + ) | ||
| 352 | return [] | 354 | return [] |
| 353 | if not isinstance(content, dict): | 355 | if not isinstance(content, dict): |
| 354 | logger.warning(f"Invalid custom api yaml format: {yaml_path}, expected dict") | 356 | logger.warning(f"Invalid custom api yaml format: {yaml_path}, expected dict") |
| @@ -30,21 +30,29 @@ from msprobe.pytorch.common.log import logger | |||
| 30 | from msprobe.pytorch.common.utils import is_torch_nn_module, register_forward_pre_hook | 30 | from msprobe.pytorch.common.utils import is_torch_nn_module, register_forward_pre_hook |
| 31 | from msprobe.pytorch.dump.module_dump.hook_wrapper import ( | 31 | from msprobe.pytorch.dump.module_dump.hook_wrapper import ( |
| 32 | wrap_setup_input_output_hook, | 32 | wrap_setup_input_output_hook, |
| 33 | - wrap_backward_hook_function_apply | 33 | + wrap_backward_hook_function_apply, |
| 34 | ) | 34 | ) |
| 35 | 35 | ||
| 36 | 36 | ||
| 37 | torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0' | 37 | torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0' |
| 38 | torch_version_above_or_equal_21 = torch.__version__.split('+')[0] >= '2.1' | 38 | torch_version_above_or_equal_21 = torch.__version__.split('+')[0] >= '2.1' |
| 39 | if torch_version_above_or_equal_21: | 39 | if torch_version_above_or_equal_21: |
| 40 | - from torch.utils.checkpoint import _StopRecomputationError | 40 | + from torch.utils.checkpoint import ( # pylint: disable=ungrouped-imports |
| 41 | + _StopRecomputationError, | ||
 | |||
| 42 | + ) | ||
| 43 | +else: | ||
| 44 | + _StopRecomputationError = type('_StopRecomputationError', (Exception,), {}) | ||
| 41 | 45 | ||
| 42 | 46 | ||
| 43 | def wrap_megatron_deallocate(func): | 47 | def wrap_megatron_deallocate(func): |
| 44 | def wrapper_func(out, deallocate_pipeline_outputs=False): | 48 | def wrapper_func(out, deallocate_pipeline_outputs=False): |
| 45 | if deallocate_pipeline_outputs and isinstance(out, torch.Tensor) and getattr(out, "_base") is not None: | 49 | if deallocate_pipeline_outputs and isinstance(out, torch.Tensor) and getattr(out, "_base") is not None: |
| 46 | out_clone = out.clone() | 50 | out_clone = out.clone() |
| 47 | - out.data = torch.empty((1,), device=out.device, dtype=out.dtype, ) | 51 | + out.data = torch.empty( |
| 52 | + (1,), | ||
| 53 | + device=out.device, | ||
| 54 | + dtype=out.dtype, | ||
| 55 | + ) | ||
| 48 | return func(out_clone, deallocate_pipeline_outputs) | 56 | return func(out_clone, deallocate_pipeline_outputs) |
| 49 | return func(out, deallocate_pipeline_outputs) | 57 | return func(out, deallocate_pipeline_outputs) |
| 50 | 58 | ||
| @@ -92,6 +100,7 @@ class ModuleProcessor: | |||
| 92 | wrap_backward_hook_function_apply() | 100 | wrap_backward_hook_function_apply() |
| 93 | try: | 101 | try: |
| 94 | from megatron.core.pipeline_parallel import schedules | 102 | from megatron.core.pipeline_parallel import schedules |
| 103 | + | ||
| 95 | origin_func_id = id(schedules.deallocate_output_tensor) | 104 | origin_func_id = id(schedules.deallocate_output_tensor) |
| 96 | schedules.deallocate_output_tensor = wrap_megatron_deallocate(schedules.deallocate_output_tensor) | 105 | schedules.deallocate_output_tensor = wrap_megatron_deallocate(schedules.deallocate_output_tensor) |
| 97 | schedules.forward_step = wrap_megatron_step(schedules.forward_step) | 106 | schedules.forward_step = wrap_megatron_step(schedules.forward_step) |
| @@ -128,9 +137,11 @@ class ModuleProcessor: | |||
| 128 | 137 | ||
| 129 | @staticmethod | 138 | @staticmethod |
| 130 | def has_register_backward_hook(module): | 139 | def has_register_backward_hook(module): |
| 131 | - return hasattr(module, '_backward_hooks') and \ | 140 | + return ( |
| 132 | - len(module._backward_hooks) > 0 and \ | 141 | + hasattr(module, '_backward_hooks') |
| 133 | - module._is_full_backward_hook is False | 142 | + and len(module._backward_hooks) > 0 |
| 143 | + and module._is_full_backward_hook is False | ||
| 144 | + ) | ||
| 134 | 145 | ||
| 135 | @staticmethod | 146 | @staticmethod |
| 136 | def get_modules_and_names(models, recursive, module_names): | 147 | def get_modules_and_names(models, recursive, module_names): |
| @@ -139,13 +150,13 @@ class ModuleProcessor: | |||
| 139 | if not recursive and len(module_names) != len(models): | 150 | if not recursive and len(module_names) != len(models): |
| 140 | return modules_and_names_with_index | 151 | return modules_and_names_with_index |
| 141 | for index, model in enumerate(models): | 152 | for index, model in enumerate(models): |
| 142 | - modules_and_names_with_index[str(index)] = model.named_modules() if recursive else \ | 153 | + modules_and_names_with_index[str(index)] = ( |
| 143 | - [(module_names[index], model)] | 154 | + model.named_modules() if recursive else [(module_names[index], model)] |
| 155 | + ) | ||
| 144 | else: | 156 | else: |
| 145 | if not recursive and len(module_names) != 1: | 157 | if not recursive and len(module_names) != 1: |
| 146 | return modules_and_names_with_index | 158 | return modules_and_names_with_index |
| 147 | - modules_and_names_with_index["-1"] = models.named_modules() if recursive else \ | 159 | + modules_and_names_with_index["-1"] = models.named_modules() if recursive else [(module_names[0], models)] |
| 148 | - [(module_names[0], models)] | ||
| 149 | return modules_and_names_with_index | 160 | return modules_and_names_with_index |
| 150 | 161 | ||
| 151 | @staticmethod | 162 | @staticmethod |
| @@ -154,8 +165,8 @@ class ModuleProcessor: | |||
| 154 | for model in models: | 165 | for model in models: |
| 155 | if "megatron" in type(model).__module__: | 166 | if "megatron" in type(model).__module__: |
| 156 | return True | 167 | return True |
| 157 | - else: | 168 | + return False |
| 158 | - return "megatron" in type(models).__module__ | 169 | + return "megatron" in type(models).__module__ |
| 159 | 170 | ||
| 160 | @classmethod | 171 | @classmethod |
| 161 | def reset_module_stats(cls): | 172 | def reset_module_stats(cls): |
| @@ -175,11 +186,11 @@ class ModuleProcessor: | |||
| 175 | modules_and_names_with_index = self.get_modules_and_names(models, recursive, module_names) | 186 | modules_and_names_with_index = self.get_modules_and_names(models, recursive, module_names) |
| 176 | for index, modules_and_names in modules_and_names_with_index.items(): | 187 | for index, modules_and_names in modules_and_names_with_index.items(): |
| 177 | model = models if index == "-1" else models[int(index)] | 188 | model = models if index == "-1" else models[int(index)] |
| 178 | - | 189 | + |
| 179 | model_list = [] | 190 | model_list = [] |
| 180 | for name, module in modules_and_names: | 191 | for name, module in modules_and_names: |
| 181 | model_list.append((name, module)) | 192 | model_list.append((name, module)) |
| 182 | - | 193 | + |
| 183 | is_verl = "verl" in sys.modules | 194 | is_verl = "verl" in sys.modules |
| 184 | for idx, (name, module) in enumerate(model_list): | 195 | for idx, (name, module) in enumerate(model_list): |
| 185 | if recursive and module == model: | 196 | if recursive and module == model: |
| @@ -192,17 +203,21 @@ class ModuleProcessor: | |||
| 192 | continue | 203 | continue |
| 193 | if module.__class__.__name__ == "FullyShardedDataParallel": | 204 | if module.__class__.__name__ == "FullyShardedDataParallel": |
| 194 | continue | 205 | continue |
| 195 | - | 206 | + |
| 196 | # verl 场景下跳过第一层和最后一层 | 207 | # verl 场景下跳过第一层和最后一层 |
| 197 | - if is_verl and (idx == 1 or idx == len(model_list) - 1): | 208 | + if is_verl and idx in (1, len(model_list) - 1): |
| 198 | - logger.warning(f"The module {name} is the first or last layer in verl scenario, " | 209 | + logger.warning( |
| 199 | - f"the data dump for this module will be skipped.") | 210 | + f"The module {name} is the first or last layer in verl scenario, " |
| 211 | + f"the data dump for this module will be skipped." | ||
| 212 | + ) | ||
| 200 | continue | 213 | continue |
| 201 | - | 214 | + |
| 202 | setattr(module, 'msprobe_hook', True) | 215 | setattr(module, 'msprobe_hook', True) |
| 203 | module_index = (index + Const.SEP) if index != "-1" else "" | 216 | module_index = (index + Const.SEP) if index != "-1" else "" |
| 204 | - prefix_name = f'{BaseScope.Module_Type_Module}{Const.SEP}{module_index}{name}{Const.SEP}' + \ | 217 | + prefix_name = ( |
| 205 | - f'{module.__class__.__name__}{Const.SEP}' | 218 | + f'{BaseScope.Module_Type_Module}{Const.SEP}{module_index}{name}{Const.SEP}' |
| 219 | + + f'{module.__class__.__name__}{Const.SEP}' | ||
| 220 | + ) | ||
| 206 | 221 | ||
| 207 | forward_pre_hook = self.build_module_hook(prefix_name, build_hook) | 222 | forward_pre_hook = self.build_module_hook(prefix_name, build_hook) |
| 208 | 223 | ||
| @@ -231,6 +246,7 @@ class ModuleProcessor: | |||
| 231 | full_forward_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}' | 246 | full_forward_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}' |
| 232 | full_backward_name = f'{module_name}{Const.BACKWARD}{Const.SEP}{index}' | 247 | full_backward_name = f'{module_name}{Const.BACKWARD}{Const.SEP}{index}' |
| 233 | 248 | ||
| 249 | + logger.debug(f"module forward_pre_hook executed for {module_name}") | ||
| 234 | self.set_construct_info_in_pre_hook(full_forward_name) | 250 | self.set_construct_info_in_pre_hook(full_forward_name) |
| 235 | 251 | ||
| 236 | if not hasattr(module, 'msprobe_forward_hook'): | 252 | if not hasattr(module, 'msprobe_forward_hook'): |
| @@ -256,6 +272,7 @@ class ModuleProcessor: | |||
| 256 | def get_backward_hook(backward_data_hook, full_backward_name): | 272 | def get_backward_hook(backward_data_hook, full_backward_name): |
| 257 | @ThreadSafe.synchronized | 273 | @ThreadSafe.synchronized |
| 258 | def backward_hook_fn(module, grad_input, grad_output): | 274 | def backward_hook_fn(module, grad_input, grad_output): |
| 275 | + logger.debug(f"module backward_hook executed for {full_backward_name}") | ||
| 259 | new_output = backward_data_hook(module, grad_input, grad_output) | 276 | new_output = backward_data_hook(module, grad_input, grad_output) |
| 260 | self.set_construct_info_in_hook(full_backward_name, is_forward=False) | 277 | self.set_construct_info_in_hook(full_backward_name, is_forward=False) |
| 261 | return new_output | 278 | return new_output |
| @@ -281,6 +298,7 @@ class ModuleProcessor: | |||
| 281 | index = ModuleProcessor.module_count.get(module_name) | 298 | index = ModuleProcessor.module_count.get(module_name) |
| 282 | full_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}' | 299 | full_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}' |
| 283 | 300 | ||
| 301 | + logger.debug(f"module forward_hook executed for {module_name}") | ||
| 284 | hook_set = build_data_hook(BaseScope.Module_Type_Module, full_name) | 302 | hook_set = build_data_hook(BaseScope.Module_Type_Module, full_name) |
| 285 | hook_result = hook_set.forward_hook(module, args, kwargs_or_output, output_or_kwargs) | 303 | hook_result = hook_set.forward_hook(module, args, kwargs_or_output, output_or_kwargs) |
| 286 | self.set_construct_info_in_hook(full_name) | 304 | self.set_construct_info_in_hook(full_name) |
| @@ -304,12 +322,14 @@ class ModuleProcessor: | |||
| 304 | ModuleProcessor.module_stack[tid] = [] | 322 | ModuleProcessor.module_stack[tid] = [] |
| 305 | 323 | ||
| 306 | if self.module_stack[tid]: | 324 | if self.module_stack[tid]: |
| 307 | - ModuleProcessor.module_node[full_name] = self.module_stack[tid][-1] if not is_megatron() \ | 325 | + ModuleProcessor.module_node[full_name] = ( |
| 308 | - else [self.module_stack[tid][-1], get_micro_step()] | 326 | + self.module_stack[tid][-1] if not is_megatron() else [self.module_stack[tid][-1], get_micro_step()] |
| 327 | + ) | ||
| 309 | else: | 328 | else: |
| 310 | parent_name = ModuleProcessor.module_queue.find_last(full_name) | 329 | parent_name = ModuleProcessor.module_queue.find_last(full_name) |
| 311 | - ModuleProcessor.module_node[full_name] = parent_name if not is_megatron() \ | 330 | + ModuleProcessor.module_node[full_name] = ( |
| 312 | - else [parent_name, get_micro_step()] | 331 | + parent_name if not is_megatron() else [parent_name, get_micro_step()] |
| 332 | + ) | ||
| 313 | 333 | ||
| 314 | ModuleProcessor.module_queue.add_name(full_name) | 334 | ModuleProcessor.module_queue.add_name(full_name) |
| 315 | ModuleProcessor.module_stack[tid].append(full_name) | 335 | ModuleProcessor.module_stack[tid].append(full_name) |
| @@ -325,8 +345,11 @@ class ModuleProcessor: | |||
| 325 | if self.module_stack.get(tid): | 345 | if self.module_stack.get(tid): |
| 326 | ModuleProcessor.module_stack[tid].pop() | 346 | ModuleProcessor.module_stack[tid].pop() |
| 327 | if self.module_stack.get(tid): | 347 | if self.module_stack.get(tid): |
| 328 | - ModuleProcessor.api_parent_node[tid] = ModuleProcessor.module_stack[tid][-1] if not is_megatron() \ | 348 | + ModuleProcessor.api_parent_node[tid] = ( |
| 349 | + ModuleProcessor.module_stack[tid][-1] | ||
| 350 | + if not is_megatron() | ||
| 329 | else [ModuleProcessor.module_stack[tid][-1], get_micro_step()] | 351 | else [ModuleProcessor.module_stack[tid][-1], get_micro_step()] |
| 352 | + ) | ||
| 330 | if self.scope: | 353 | if self.scope: |
| 331 | self.scope.end_module(full_name) | 354 | self.scope.end_module(full_name) |
| 332 | else: | 355 | else: |
| @@ -23,7 +23,7 @@ from msprobe.pytorch.dump.api_dump.api_register import ( | |||
| 23 | get_api_register, | 23 | get_api_register, |
| 24 | ApiTemplate, | 24 | ApiTemplate, |
| 25 | redirect_wait, | 25 | redirect_wait, |
| 26 | - reset_dist_collect_func | 26 | + reset_dist_collect_func, |
| 27 | ) | 27 | ) |
| 28 | from msprobe.pytorch.dump.api_dump.hook_module import HOOKModule | 28 | from msprobe.pytorch.dump.api_dump.hook_module import HOOKModule |
| 29 | from msprobe.pytorch.dump.api_dump.pt_hook_manager import PytorchHookManager | 29 | from msprobe.pytorch.dump.api_dump.pt_hook_manager import PytorchHookManager |
| @@ -54,12 +54,13 @@ class PytorchService(BaseService): | |||
| 54 | def _init_specific_components(self): | 54 | def _init_specific_components(self): |
| 55 | self.logger = logger | 55 | self.logger = logger |
| 56 | self.api_register = get_api_register() | 56 | self.api_register = get_api_register() |
| 57 | + self.module_processor = None | ||
| 57 | self._refresh_module_processor() | 58 | self._refresh_module_processor() |
| 58 | self.hook_manager = PytorchHookManager(self.data_collector, self.config) | 59 | self.hook_manager = PytorchHookManager(self.data_collector, self.config) |
| 59 | self.api_template = ApiTemplate | 60 | self.api_template = ApiTemplate |
| 60 | 61 | ||
| 61 | def _refresh_module_processor(self): | 62 | def _refresh_module_processor(self): |
| 62 | - self.module_processor = ModuleProcessor(self.data_collector.scope) | 63 | + self.module_processor = ModuleProcessor(self.data_collector.scope) # pylint: disable=attribute-defined-outside-init |
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。   | |||
| 63 | 64 | ||
| 64 | def _register_hook(self): | 65 | def _register_hook(self): |
| 65 | if self._is_mix_level: | 66 | if self._is_mix_level: |
| @@ -95,7 +96,7 @@ class PytorchService(BaseService): | |||
| 95 | def _register_module_hook(self): | 96 | def _register_module_hook(self): |
| 96 | ModuleProcessor.enable_module_dump = True | 97 | ModuleProcessor.enable_module_dump = True |
| 97 | self.module_processor.register_module_hook(self.model, self.build_hook) | 98 | self.module_processor.register_module_hook(self.model, self.build_hook) |
| 98 | - self.logger.info(f"The module {self.config.task} hook function is successfully mounted to the model.") | 99 | + self.logger.info_on_rank_0(f"The module {self.config.task} hook function is successfully mounted to the model.") |
| 99 | 100 | ||
| 100 | def _reset_status(self): | 101 | def _reset_status(self): |
| 101 | super()._reset_status() | 102 | super()._reset_status() |
| @@ -136,7 +136,7 @@ class TestBaseService(unittest.TestCase): | |||
| 136 | mock_register_hook.assert_called_once() | 136 | mock_register_hook.assert_called_once() |
| 137 | mock_register_module_hook.assert_called_once() | 137 | mock_register_module_hook.assert_called_once() |
| 138 | 138 | ||
| 139 | - self.service.logger.info.assert_called_with(f"Dump data will be saved in {self.service.dump_iter_dir}.") | 139 | + self.service.logger.info_on_rank_0.assert_called_with(f"Dump data will be saved in {self.service.dump_iter_dir}.") |
| 140 | self.assertTrue(Runtime.is_running) | 140 | self.assertTrue(Runtime.is_running) |
| 141 | self.assertTrue(self.service.primitive_switch) | 141 | self.assertTrue(self.service.primitive_switch) |
| 142 | self.assertFalse(self.service.first_start) | 142 | self.assertFalse(self.service.first_start) |
| @@ -186,7 +186,7 @@ class TestBaseService(unittest.TestCase): | |||
| 186 | self.assertFalse(Runtime.is_running) | 186 | self.assertFalse(Runtime.is_running) |
| 187 | self.assertFalse(self.service.primitive_switch) | 187 | self.assertFalse(self.service.primitive_switch) |
| 188 | 188 | ||
| 189 | - self.service.logger.info.assert_called_with( | 189 | + self.service.logger.info_on_rank_0.assert_called_with( |
| 190 | f"{Const.TOOL_NAME}: debugger.stop() is set successfully. " | 190 | f"{Const.TOOL_NAME}: debugger.stop() is set successfully. " |
| 191 | "Please set debugger.start() to turn on the dump switch again. " | 191 | "Please set debugger.start() to turn on the dump switch again. " |
| 192 | ) | 192 | ) |
| @@ -329,7 +329,7 @@ class TestBaseService(unittest.TestCase): | |||
| 329 | self.service._register_api_hook() | 329 | self.service._register_api_hook() |
| 330 | self.service.api_register.initialize_hook.assert_called() | 330 | self.service.api_register.initialize_hook.assert_called() |
| 331 | self.service.api_register.register_all_api.assert_called() | 331 | self.service.api_register.register_all_api.assert_called() |
| 332 | - self.service.logger.info.assert_called_with( | 332 | + self.service.logger.info_on_rank_0.assert_called_with( |
| 333 | f"The api {self.config.task} hook function is successfully mounted to the model." | 333 | f"The api {self.config.task} hook function is successfully mounted to the model." |
| 334 | ) | 334 | ) |
| 335 | 335 | ||
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。