已合并
feat(monitor): add debug logs for issue 292 #814
qiange_star创建于 6月22日
feat(monitor): add debug logs for issue 292 #814
已合并
共 4 个文件变更+303-20
| @@ -352,6 +352,16 @@ class TrainerMon: | |||
| 352 | self.cc_pre_hook = self.cc_distribution.get('cc_pre_hook', False) | 352 | self.cc_pre_hook = self.cc_distribution.get('cc_pre_hook', False) |
| 353 | 353 | ||
| 354 | self.common_info() | 354 | self.common_info() |
| 355 | + logger.debug( | ||
| 356 | + "set_config normalized: " | ||
| 357 | + f"start_step={self.start_step}, collect_times={self.collect_times}, " | ||
| 358 | + f"step_interval={self.step_interval}, format={self.format}, ops_count={len(self.ops)}, " | ||
| 359 | + f"module_ranks={self.module_rank_list}, xy_distribution={self.xy_distribution}, " | ||
| 360 | + f"wg_distribution={self.wg_distribution}, mv_distribution={self.mv_distribution}, " | ||
| 361 | + f"param_distribution={self.param_distribution}, cc_enabled={self.cc_distribution.get('enable', False)}, " | ||
| 362 | + f"monitor_mbs_grad={self.monitor_mbs_grad}, print_struct={self.print_struct}, " | ||
| 363 | + f"stack_info={self.stack_info}, recording_l2_features={self.recording_l2_features}" | ||
| 364 | + ) | ||
| 355 | 365 | ||
| 356 | # 初始化AnomalyData工厂 | 366 | # 初始化AnomalyData工厂 |
| 357 | alert_setting = self.config.get('alert', {"rules": []}) | 367 | alert_setting = self.config.get('alert', {"rules": []}) |
| @@ -371,8 +381,14 @@ class TrainerMon: | |||
| 371 | 381 | ||
| 372 | writer = FORMAT_MAPPING[self.format] | 382 | writer = FORMAT_MAPPING[self.format] |
| 373 | self.step_count_per_record = self.config.get('step_count_per_record', 1) | 383 | self.step_count_per_record = self.config.get('step_count_per_record', 1) |
| 384 | + should_init_writer = (self.rank in self.module_rank_list) or len(self.module_rank_list) == 0 | ||
| 385 | + logger.debug( | ||
| 386 | + f"set_config writer setup: writer={writer.__name__}, should_init_writer={should_init_writer}, " | ||
| 387 | + f"tensorboard_dir={self.tensorboard_dir}, step_count_per_record={self.step_count_per_record}, " | ||
| 388 | + f"anomaly_dump={self.anomaly_data_factory is not None}" | ||
| 389 | + ) | ||
| 374 | 390 | ||
| 375 | - if (self.rank in self.module_rank_list) or len(self.module_rank_list) == 0: | 391 | + if should_init_writer: |
| 376 | self.summary_writer = writer( | 392 | self.summary_writer = writer( |
| 377 | WriterInput( | 393 | WriterInput( |
| 378 | self.tensorboard_dir, | 394 | self.tensorboard_dir, |
| @@ -516,9 +532,18 @@ class TrainerMon: | |||
| 516 | 532 | ||
| 517 | def generate_wgrad_metrics(self, post_grad_dict): | 533 | def generate_wgrad_metrics(self, post_grad_dict): |
| 518 | if not self.wg_distribution: | 534 | if not self.wg_distribution: |
| 535 | + logger.debug("generate_wgrad_metrics skipped because wg_distribution is disabled") | ||
| 519 | return {}, {} | 536 | return {}, {} |
| 520 | 537 | ||
| 538 | + logger.debug( | ||
| 539 | + f"generate_wgrad_metrics start: post_grad_count={len(post_grad_dict)}, " | ||
| 540 | + f"existing_post_metric_count={len(self.grad_context.post)}" | ||
| 541 | + ) | ||
| 521 | get_metrics(self.ops, post_grad_dict, self.eps, self.grad_context.post) | 542 | get_metrics(self.ops, post_grad_dict, self.eps, self.grad_context.post) |
| 543 | + logger.debug( | ||
| 544 | + f"generate_wgrad_metrics finish: post_metric_count={len(self.grad_context.post)}, " | ||
| 545 | + f"pre_metric_count={len(self.grad_context.pre)}" | ||
| 546 | + ) | ||
| 522 | return self.grad_context.post, self.grad_context.pre | 547 | return self.grad_context.post, self.grad_context.pre |
| 523 | 548 | ||
| 524 | def generate_xy_metrics(self): | 549 | def generate_xy_metrics(self): |
| @@ -604,8 +629,13 @@ class TrainerMon: | |||
| 604 | 629 | ||
| 605 | def write_grad_tb(self, step): | 630 | def write_grad_tb(self, step): |
| 606 | if not self.wg_distribution: | 631 | if not self.wg_distribution: |
| 632 | + logger.debug(f"write_grad_tb skipped at step={step} because wg_distribution is disabled") | ||
| 607 | return | 633 | return |
| 608 | 634 | ||
| 635 | + logger.debug( | ||
| 636 | + f"write_grad_tb start: step={step}, monitor_mbs_grad={self.monitor_mbs_grad}, " | ||
| 637 | + f"pre_metric_count={len(self.grad_context.pre)}, post_metric_count={len(self.grad_context.post)}" | ||
| 638 | + ) | ||
| 609 | self.summary_writer.write_metrics( | 639 | self.summary_writer.write_metrics( |
| 610 | self.ops, self.grad_context.pre, step, 'grad_unreduced', use_micro_step=self.monitor_mbs_grad | 640 | self.ops, self.grad_context.pre, step, 'grad_unreduced', use_micro_step=self.monitor_mbs_grad |
| 611 | ) | 641 | ) |
| @@ -615,6 +645,12 @@ class TrainerMon: | |||
| 615 | # in DDP by default use params_have_main_grad | 645 | # in DDP by default use params_have_main_grad |
| 616 | def optimizer_pre_step_hook(optimizer, args, kwargs): | 646 | def optimizer_pre_step_hook(optimizer, args, kwargs): |
| 617 | context = self.optimizer_context[optimizer] | 647 | context = self.optimizer_context[optimizer] |
| 648 | + logger.debug( | ||
| 649 | + f"optimizer_pre_step_hook enter: step={context.step}, start_step={self.start_step}, " | ||
| 650 | + f"step_interval={self.step_interval}, optimizer_type={optimizer.__class__.__name__}, " | ||
| 651 | + f"wg_distribution={self.wg_distribution}, mv_distribution={self.mv_distribution}, " | ||
| 652 | + f"ur_distribution={self.ur_distribution}, mg_direction={self.mg_direction}" | ||
| 653 | + ) | ||
| 618 | 654 | ||
| 619 | if ( | 655 | if ( |
| 620 | self.print_struct | 656 | self.print_struct |
| @@ -633,11 +669,19 @@ class TrainerMon: | |||
| 633 | 669 | ||
| 634 | # skip generate metrics | 670 | # skip generate metrics |
| 635 | if context.step < self.start_step or (context.step - self.start_step) % self.step_interval != 0: | 671 | if context.step < self.start_step or (context.step - self.start_step) % self.step_interval != 0: |
| 672 | + logger.debug( | ||
| 673 | + f"optimizer_pre_step_hook skip: step={context.step}, start_step={self.start_step}, " | ||
| 674 | + f"step_interval={self.step_interval}" | ||
| 675 | + ) | ||
| 636 | return | 676 | return |
| 637 | 677 | ||
| 638 | grad_dict = {} | 678 | grad_dict = {} |
| 639 | if self.wg_distribution: | 679 | if self.wg_distribution: |
| 640 | grad_dict = self.optimizer_mon.fetch_grad(self, self.param2name) | 680 | grad_dict = self.optimizer_mon.fetch_grad(self, self.param2name) |
| 681 | + logger.debug( | ||
| 682 | + f"optimizer_pre_step_hook fetched grads: grad_count={len(grad_dict)}, " | ||
| 683 | + f"target_param_count={len(self.param2name)}" | ||
| 684 | + ) | ||
| 641 | 685 | ||
| 642 | mv_result = None | 686 | mv_result = None |
| 643 | if self.mv_distribution or self.ur_distribution or self.mg_direction: | 687 | if self.mv_distribution or self.ur_distribution or self.mg_direction: |
| @@ -647,10 +691,21 @@ class TrainerMon: | |||
| 647 | context.param_exp_avg_sq = mv_result.exp_avg_sq | 691 | context.param_exp_avg_sq = mv_result.exp_avg_sq |
| 648 | context.param_adam_update = mv_result.update | 692 | context.param_adam_update = mv_result.update |
| 649 | context.param_adam_ratio = mv_result.ratio | 693 | context.param_adam_ratio = mv_result.ratio |
| 694 | + logger.debug( | ||
| 695 | + f"optimizer_pre_step_hook fetched mv: exp_avg_count={len(mv_result.exp_avg)}, " | ||
| 696 | + f"exp_avg_sq_count={len(mv_result.exp_avg_sq)}, update_count={len(mv_result.update)}, " | ||
| 697 | + f"ratio_count={len(mv_result.ratio)}" | ||
| 698 | + ) | ||
| 650 | 699 | ||
| 651 | _, _ = self.generate_wgrad_metrics(grad_dict) | 700 | _, _ = self.generate_wgrad_metrics(grad_dict) |
| 652 | self.generate_mv_metrics(context) | 701 | self.generate_mv_metrics(context) |
| 653 | self.generate_param_metrics(context, MonitorConst.PRE_PARAM) | 702 | self.generate_param_metrics(context, MonitorConst.PRE_PARAM) |
| 703 | + logger.debug( | ||
| 704 | + f"optimizer_pre_step_hook metrics prepared: post_grad_metric_count={len(self.grad_context.post)}, " | ||
| 705 | + f"pre_grad_metric_count={len(self.grad_context.pre)}, param_metric_count={len(context.param_metric)}, " | ||
| 706 | + f"exp_avg_metric_count={len(context.exp_avg_metric)}, " | ||
| 707 | + f"exp_avg_sq_metric_count={len(context.exp_avg_sq_metric)}" | ||
| 708 | + ) | ||
 | |||
| 654 | 709 | ||
| 655 | tbtag_tensor_map = {} | 710 | tbtag_tensor_map = {} |
| 656 | if self.mg_direction: | 711 | if self.mg_direction: |
| @@ -673,7 +728,12 @@ class TrainerMon: | |||
| 673 | metric_dict.update(cc.data) | 728 | metric_dict.update(cc.data) |
| 674 | cc.reset() | 729 | cc.reset() |
| 675 | 730 | ||
| 731 | + logger.debug( | ||
| 732 | + f"optimizer_pre_step_hook metric aggregation: mg_direction_metric_count={len(tbtag_tensor_map)}, " | ||
| 733 | + f"other_metric_count={len(metric_dict)}" | ||
| 734 | + ) | ||
| 676 | if not metric_dict: | 735 | if not metric_dict: |
| 736 | + logger.debug(f"optimizer_pre_step_hook no extra metrics generated at step={context.step}") | ||
| 677 | return | 737 | return |
| 678 | context.metric_dict = metric_dict | 738 | context.metric_dict = metric_dict |
| 679 | return | 739 | return |
| @@ -683,8 +743,10 @@ class TrainerMon: | |||
| 683 | self.generate_param_metrics(context, MonitorConst.POST_PARAM) | 743 | self.generate_param_metrics(context, MonitorConst.POST_PARAM) |
| 684 | 744 | ||
| 685 | if self.optimizer_hooked: | 745 | if self.optimizer_hooked: |
| 746 | + logger.debug("hook_optimizer skipped because optimizer hooks are already registered") | ||
| 686 | return | 747 | return |
| 687 | 748 | ||
| 749 | + logger.debug(f"hook_optimizer register hooks for optimizer_type={optimizer.__class__.__name__}") | ||
| 688 | self.pre_step_hooks.append(optimizer_pre_step_hook) | 750 | self.pre_step_hooks.append(optimizer_pre_step_hook) |
| 689 | self.post_step_hooks.append(optimizer_post_step_hook) | 751 | self.post_step_hooks.append(optimizer_post_step_hook) |
| 690 | 752 | ||
| @@ -740,8 +802,18 @@ class TrainerMon: | |||
| 740 | step_condition = ( | 802 | step_condition = ( |
| 741 | context.step >= self.start_step and (context.step - self.start_step) % self.step_interval == 0 | 803 | context.step >= self.start_step and (context.step - self.start_step) % self.step_interval == 0 |
| 742 | ) | 804 | ) |
| 805 | + logger.debug( | ||
| 806 | + f"step_final_hook enter: step={context.step}, monitoring={self.monitoring}, rank={rank}, " | ||
| 807 | + f"module_rank_valid={module_rank_valid}, step_condition={step_condition}, " | ||
| 808 | + f"has_collect_times={self.has_collect_times}, collect_times={self.collect_times}" | ||
| 809 | + ) | ||
| 743 | if module_rank_valid and step_condition: | 810 | if module_rank_valid and step_condition: |
| 744 | self.has_collect_times += 1 | 811 | self.has_collect_times += 1 |
| 812 | + logger.debug( | ||
| 813 | + f"step_final_hook write path: step={context.step}, collect_progress=" | ||
| 814 | + f"{self.has_collect_times}/{self.collect_times}, pre_grad_metric_count={len(self.grad_context.pre)}, " | ||
| 815 | + f"post_grad_metric_count={len(self.grad_context.post)}, other_metric_count={len(context.metric_dict)}" | ||
| 816 | + ) | ||
| 745 | 817 | ||
| 746 | if self.anomaly_data_factory: | 818 | if self.anomaly_data_factory: |
| 747 | self.anomaly_data_factory.set_call_id(self.param_name_call_id) | 819 | self.anomaly_data_factory.set_call_id(self.param_name_call_id) |
| @@ -784,9 +856,13 @@ class TrainerMon: | |||
| 784 | self.param_name_call_id.clear() | 856 | self.param_name_call_id.clear() |
| 785 | 857 | ||
| 786 | if self.has_collect_times >= self.collect_times: | 858 | if self.has_collect_times >= self.collect_times: |
| 859 | + logger.debug( | ||
| 860 | + f"step_final_hook reached collect limit at step={context.step}, removing all hooks" | ||
| 861 | + ) | ||
| 787 | self._remove_all_hooks_final(optimizer) | 862 | self._remove_all_hooks_final(optimizer) |
| 788 | 863 | ||
| 789 | context.step += 1 | 864 | context.step += 1 |
| 865 | + logger.debug(f"step_final_hook finish: next_step={context.step}") | ||
| 790 | self.dynamic_monitor(optimizer) | 866 | self.dynamic_monitor(optimizer) |
| 791 | 867 | ||
| 792 | def patch_step(func, optimizer): | 868 | def patch_step(func, optimizer): |
| @@ -959,15 +1035,22 @@ class TrainerMon: | |||
| 959 | return False | 1035 | return False |
| 960 | 1036 | ||
| 961 | def _register_chunk(self, model_chunk, prefix): | 1037 | def _register_chunk(self, model_chunk, prefix): |
| 1038 | + logger.debug( | ||
| 1039 | + f"_register_chunk start: prefix={prefix}, model_chunk_type={model_chunk.__class__.__name__}, " | ||
| 1040 | + f"fsdp_wrapped_module={self.fsdp_wrapped_module}, fsdp2_wrapped_module={self.fsdp2_wrapped_module}" | ||
| 1041 | + ) | ||
| 962 | if isinstance(model_chunk, FSDP): | 1042 | if isinstance(model_chunk, FSDP): |
| 963 | if not model_chunk._use_orig_params: | 1043 | if not model_chunk._use_orig_params: |
| 964 | raise ValueError("Only Support fsdp1 with use_orig_params=True") | 1044 | raise ValueError("Only Support fsdp1 with use_orig_params=True") |
| 965 | self.fsdp_wrapped_module = True | 1045 | self.fsdp_wrapped_module = True |
| 1046 | + logger.debug(f"_register_chunk detected fsdp1 wrapper for prefix={prefix}") | ||
| 1047 | + registered_count = 0 | ||
| 966 | for param_name, param in model_chunk.named_parameters(): | 1048 | for param_name, param in model_chunk.named_parameters(): |
| 967 | if not param.requires_grad: | 1049 | if not param.requires_grad: |
| 968 | continue | 1050 | continue |
| 969 | if not self.fsdp2_wrapped_module and param.__class__.__name__ == "DTensor": | 1051 | if not self.fsdp2_wrapped_module and param.__class__.__name__ == "DTensor": |
| 970 | self.fsdp2_wrapped_module = True | 1052 | self.fsdp2_wrapped_module = True |
| 1053 | + logger.debug(f"_register_chunk detected fsdp2 dtensor param: param_name={param_name}, prefix={prefix}") | ||
| 971 | if self.fsdp_wrapped_module: # FSDP1需要记录完整的不被target限制的flat权重名，以供后续对flat解包 | 1054 | if self.fsdp_wrapped_module: # FSDP1需要记录完整的不被target限制的flat权重名，以供后续对flat解包 |
| 972 | self.fsdp_param_name_map[id(param)] = param_name | 1055 | self.fsdp_param_name_map[id(param)] = param_name |
| 973 | 1056 | ||
| @@ -991,6 +1074,16 @@ class TrainerMon: | |||
| 991 | MonitorConst.POST_PARAM, | 1074 | MonitorConst.POST_PARAM, |
| 992 | ] | 1075 | ] |
| 993 | self.name2tag[name] = {k: get_summary_writer_tag_name(name, k, self.rank) for k in keywords} | 1076 | self.name2tag[name] = {k: get_summary_writer_tag_name(name, k, self.rank) for k in keywords} |
| 1077 | + registered_count += 1 | ||
| 1078 | + logger.debug( | ||
| 1079 | + f"_register_chunk mapped param: origin_name={param_name}, mapped_name={name}, " | ||
| 1080 | + f"duplicate={self.duplicate_param.get(name, False)}, fsdp_flat_mapped={id(param) in self.fsdp_param_name_map}" | ||
| 1081 | + ) | ||
| 1082 | + | ||
| 1083 | + logger.debug( | ||
| 1084 | + f"_register_chunk finish: prefix={prefix}, registered_count={registered_count}, " | ||
| 1085 | + f"name2param_count={len(self.name2param)}, origin2squash_count={len(self.origin2squash)}" | ||
| 1086 | + ) | ||
| 994 | 1087 | ||
| 995 | def _register_param_name(self): | 1088 | def _register_param_name(self): |
| 996 | for vpp_stage, model_chunk in enumerate(self.model): | 1089 | for vpp_stage, model_chunk in enumerate(self.model): |
| @@ -1223,28 +1316,38 @@ class TrainerMon: | |||
| 1223 | 1316 | ||
| 1224 | def _patch_grad_sync(self): | 1317 | def _patch_grad_sync(self): |
| 1225 | if not self.wg_distribution: | 1318 | if not self.wg_distribution: |
| 1319 | + logger.debug("_patch_grad_sync skipped because wg_distribution is disabled") | ||
| 1226 | return | 1320 | return |
| 1227 | 1321 | ||
| 1228 | if self.fsdp_wrapped_module: | 1322 | if self.fsdp_wrapped_module: |
| 1323 | + logger.debug("_patch_grad_sync select fsdp1 post backward hook path") | ||
| 1229 | # patch fsdp _runtime_utils._post_backward_hook | 1324 | # patch fsdp _runtime_utils._post_backward_hook |
| 1230 | self._patch_fsdp_post_backward_hook() | 1325 | self._patch_fsdp_post_backward_hook() |
| 1231 | return | 1326 | return |
| 1232 | 1327 | ||
| 1233 | if self.fsdp2_wrapped_module: | 1328 | if self.fsdp2_wrapped_module: |
| 1329 | + logger.debug("_patch_grad_sync select fsdp2 foreach_reduce path") | ||
| 1234 | # patch fsdp2 _fully_shard._fsdp_collectives.foreach_reduce | 1330 | # patch fsdp2 _fully_shard._fsdp_collectives.foreach_reduce |
| 1235 | self._patch_fsdp2_foreach_reduce() | 1331 | self._patch_fsdp2_foreach_reduce() |
| 1236 | return | 1332 | return |
| 1237 | 1333 | ||
| 1238 | if self.monitor_mbs_grad: | 1334 | if self.monitor_mbs_grad: |
| 1335 | + logger.debug("_patch_grad_sync select param_hook path because monitor_mbs_grad is enabled") | ||
| 1239 | self._hook_weights() | 1336 | self._hook_weights() |
| 1240 | return | 1337 | return |
| 1241 | 1338 | ||
| 1339 | + logger.debug("_patch_grad_sync delegate to optimizer patch path") | ||
| 1242 | self.optimizer_mon.patch_grad_sync(self) | 1340 | self.optimizer_mon.patch_grad_sync(self) |
| 1243 | 1341 | ||
| 1244 | if self.enable_megatron or self.enable_deepspeed: | 1342 | if self.enable_megatron or self.enable_deepspeed: |
| 1343 | + logger.debug( | ||
| 1344 | + f"_patch_grad_sync using distributed optimizer path: enable_megatron={self.enable_megatron}, " | ||
| 1345 | + f"enable_deepspeed={self.enable_deepspeed}" | ||
| 1346 | + ) | ||
| 1245 | return | 1347 | return |
| 1246 | 1348 | ||
| 1247 | # default hook weights | 1349 | # default hook weights |
| 1350 | + logger.debug("_patch_grad_sync fallback to default param_hook path") | ||
| 1248 | self._hook_weights() | 1351 | self._hook_weights() |
| 1249 | 1352 | ||
| 1250 | def _patch_fsdp_post_backward_hook(self): | 1353 | def _patch_fsdp_post_backward_hook(self): |
| @@ -1270,10 +1373,18 @@ class TrainerMon: | |||
| 1270 | offsets = handle._get_flat_param_offsets() | 1373 | offsets = handle._get_flat_param_offsets() |
| 1271 | shapes = handle.flat_param._shapes | 1374 | shapes = handle.flat_param._shapes |
| 1272 | params = handle.flat_param._params | 1375 | params = handle.flat_param._params |
| 1376 | + logger.debug( | ||
| 1377 | + f"_post_backward_hook enter: micro_step={handle.flat_param.micro_step}, " | ||
| 1378 | + f"param_count={len(params)}, offsets_count={len(offsets)}, shape_count={len(shapes)}" | ||
| 1379 | + ) | ||
| 1273 | for param, (start, end), local_shape in zip(params, offsets, shapes): | 1380 | for param, (start, end), local_shape in zip(params, offsets, shapes): |
| 1274 | param_name = self.fsdp_param_name_map.get(id(param), "") | 1381 | param_name = self.fsdp_param_name_map.get(id(param), "") |
| 1275 | tag = self.name2tag.get(self.origin2squash.get(param_name, ""), {}).get(MonitorConst.PRE_GRAD) | 1382 | tag = self.name2tag.get(self.origin2squash.get(param_name, ""), {}).get(MonitorConst.PRE_GRAD) |
| 1276 | if tag is None: | 1383 | if tag is None: |
| 1384 | + logger.debug( | ||
| 1385 | + f"_post_backward_hook skip param without tag: param_name={param_name}, " | ||
| 1386 | + f"micro_step={handle.flat_param.micro_step}" | ||
| 1387 | + ) | ||
| 1277 | continue | 1388 | continue |
| 1278 | grad = handle.flat_param.grad[start : end + 1].reshape(local_shape) | 1389 | grad = handle.flat_param.grad[start : end + 1].reshape(local_shape) |
| 1279 | if self.monitor_mbs_grad: | 1390 | if self.monitor_mbs_grad: |
| @@ -1281,6 +1392,10 @@ class TrainerMon: | |||
| 1281 | grad_dict[tag] = grad | 1392 | grad_dict[tag] = grad |
| 1282 | self.register_param_call_id("_post_backward_hook", tag) | 1393 | self.register_param_call_id("_post_backward_hook", tag) |
| 1283 | get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre) | 1394 | get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre) |
| 1395 | + logger.debug( | ||
| 1396 | + f"_post_backward_hook collected pre_grad metrics: grad_count={len(grad_dict)}, " | ||
| 1397 | + f"pre_metric_count={len(self.grad_context.pre)}" | ||
| 1398 | + ) | ||
| 1284 | handle.flat_param.micro_step = (handle.flat_param.micro_step + 1) % self.micro_batch_number | 1399 | handle.flat_param.micro_step = (handle.flat_param.micro_step + 1) % self.micro_batch_number |
| 1285 | out = _post_backward_hook(state, handle, *unused) | 1400 | out = _post_backward_hook(state, handle, *unused) |
| 1286 | return out | 1401 | return out |
| @@ -1301,17 +1416,29 @@ class TrainerMon: | |||
| 1301 | 1416 | ||
| 1302 | if self.monitor_mbs_grad or (fsdp_params[0].micro_step + 1 == self.micro_batch_number): | 1417 | if self.monitor_mbs_grad or (fsdp_params[0].micro_step + 1 == self.micro_batch_number): |
| 1303 | grad_dict = {} | 1418 | grad_dict = {} |
| 1419 | + logger.debug( | ||
| 1420 | + f"foreach_reduce enter: micro_step={fsdp_params[0].micro_step}, " | ||
| 1421 | + f"param_count={len(fsdp_params)}, grad_count={len(unsharded_grads)}" | ||
| 1422 | + ) | ||
| 1304 | for param, grad in zip(fsdp_params, unsharded_grads): | 1423 | for param, grad in zip(fsdp_params, unsharded_grads): |
| 1305 | tag = self.name2tag.get(self.origin2squash.get(param._param_fqn, ""), {}).get( | 1424 | tag = self.name2tag.get(self.origin2squash.get(param._param_fqn, ""), {}).get( |
| 1306 | MonitorConst.PRE_GRAD | 1425 | MonitorConst.PRE_GRAD |
| 1307 | ) | 1426 | ) |
| 1308 | if tag is None: | 1427 | if tag is None: |
| 1428 | + logger.debug( | ||
| 1429 | + f"foreach_reduce skip param without tag: param_fqn={param._param_fqn}, " | ||
| 1430 | + f"micro_step={fsdp_params[0].micro_step}" | ||
| 1431 | + ) | ||
| 1309 | continue | 1432 | continue |
| 1310 | if self.monitor_mbs_grad: | 1433 | if self.monitor_mbs_grad: |
| 1311 | tag = tag.replace('/', f'{MonitorConst.NAME_SEP}{fsdp_params[0].micro_step}/', 1) | 1434 | tag = tag.replace('/', f'{MonitorConst.NAME_SEP}{fsdp_params[0].micro_step}/', 1) |
| 1312 | grad_dict[tag] = grad | 1435 | grad_dict[tag] = grad |
| 1313 | self.register_param_call_id("foreach_reduce", tag) | 1436 | self.register_param_call_id("foreach_reduce", tag) |
| 1314 | get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre) | 1437 | get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre) |
| 1438 | + logger.debug( | ||
| 1439 | + f"foreach_reduce collected pre_grad metrics: grad_count={len(grad_dict)}, " | ||
| 1440 | + f"pre_metric_count={len(self.grad_context.pre)}" | ||
| 1441 | + ) | ||
| 1315 | fsdp_params[0].micro_step = (fsdp_params[0].micro_step + 1) % self.micro_batch_number | 1442 | fsdp_params[0].micro_step = (fsdp_params[0].micro_step + 1) % self.micro_batch_number |
| 1316 | out = foreach_reduce(fsdp_params, unsharded_grads, *unused) | 1443 | out = foreach_reduce(fsdp_params, unsharded_grads, *unused) |
| 1317 | return out | 1444 | return out |
| @@ -1338,6 +1465,10 @@ class TrainerMon: | |||
| 1338 | key += f'{MonitorConst.NAME_SEP}{param.micro_step}' | 1465 | key += f'{MonitorConst.NAME_SEP}{param.micro_step}' |
| 1339 | 1466 | ||
| 1340 | key = get_summary_writer_tag_name(key, 'acc_grad', self.rank) | 1467 | key = get_summary_writer_tag_name(key, 'acc_grad', self.rank) |
| 1468 | + logger.debug( | ||
| 1469 | + f"param_hook enter: name={name}, key={key}, micro_step={param.micro_step}, " | ||
| 1470 | + f"use_main_grad={self.params_have_main_grad}" | ||
| 1471 | + ) | ||
| 1341 | self.register_param_call_id("param_hook", key) | 1472 | self.register_param_call_id("param_hook", key) |
| 1342 | if self.monitor_mbs_grad or (param.micro_step + 1 == self.micro_batch_number): | 1473 | if self.monitor_mbs_grad or (param.micro_step + 1 == self.micro_batch_number): |
| 1343 | if self.params_have_main_grad: | 1474 | if self.params_have_main_grad: |
| @@ -1345,9 +1476,11 @@ class TrainerMon: | |||
| 1345 | else: | 1476 | else: |
| 1346 | grad = param.grad | 1477 | grad = param.grad |
| 1347 | get_metrics(self.ops, {key: grad.clone()}, self.eps, self.grad_context.pre) | 1478 | get_metrics(self.ops, {key: grad.clone()}, self.eps, self.grad_context.pre) |
| 1479 | + logger.debug(f"param_hook collected metric: key={key}, pre_metric_count={len(self.grad_context.pre)}") | ||
| 1348 | param.micro_step = (param.micro_step + 1) % self.micro_batch_number | 1480 | param.micro_step = (param.micro_step + 1) % self.micro_batch_number |
| 1349 | 1481 | ||
| 1350 | logger.info("hooking weight grads.") | 1482 | logger.info("hooking weight grads.") |
| 1483 | + logger.debug(f"_hook_weights register start: param_count={len(self.param2name)}") | ||
| 1351 | for param, name in self.param2name.items(): | 1484 | for param, name in self.param2name.items(): |
| 1352 | setattr(param, 'micro_step', 0) | 1485 | setattr(param, 'micro_step', 0) |
| 1353 | param_tmp = param.expand_as(param) | 1486 | param_tmp = param.expand_as(param) |
| @@ -24,7 +24,7 @@ from msprobe.pytorch.monitor.module_metric import get_metrics | |||
| 24 | from msprobe.core.common.const import MonitorConst | 24 | from msprobe.core.common.const import MonitorConst |
| 25 | 25 | ||
| 26 | 26 | ||
| 27 | -class OptimizerMon(object): | 27 | +class OptimizerMon: |
| 28 | def __init__(self, torch_opt) -> None: | 28 | def __init__(self, torch_opt) -> None: |
| 29 | self.fp16_to_fp32_param = {} | 29 | self.fp16_to_fp32_param = {} |
| 30 | self.torch_opt = torch_opt | 30 | self.torch_opt = torch_opt |
| @@ -131,6 +131,7 @@ class OptimizerMon(object): | |||
| 131 | def patch_sync(sync_grad_func): | 131 | def patch_sync(sync_grad_func): |
| 132 | def wrapper(bucket, *args, **kwargs): | 132 | def wrapper(bucket, *args, **kwargs): |
| 133 | grad_dict = {} | 133 | grad_dict = {} |
| 134 | + param_source = "params" | ||
| 134 | # Megatron between core_r0.6.0 and core_r0.8.0, this bucket is Bucket. | 135 | # Megatron between core_r0.6.0 and core_r0.8.0, this bucket is Bucket. |
| 135 | # When megatron is core_r0.9.0, this bucket is _ParamAndGradBucketGroup. | 136 | # When megatron is core_r0.9.0, this bucket is _ParamAndGradBucketGroup. |
| 136 | # In megatron version core_r0.9.0, func start_grad_sync from Bucket moved to _ParamAndGradBucketGroup. | 137 | # In megatron version core_r0.9.0, func start_grad_sync from Bucket moved to _ParamAndGradBucketGroup. |
| @@ -138,9 +139,16 @@ class OptimizerMon(object): | |||
| 138 | param_attr = bucket.params | 139 | param_attr = bucket.params |
| 139 | elif hasattr(bucket, "params_list"): | 140 | elif hasattr(bucket, "params_list"): |
| 140 | param_attr = bucket.params_list | 141 | param_attr = bucket.params_list |
| 142 | + param_source = "params_list" | ||
| 141 | else: | 143 | else: |
| 142 | raise AttributeError( | 144 | raise AttributeError( |
| 143 | - f"class {bucket.__class__.__name__} has no attribute 'params' or 'params_list', please check!") | 145 | + f"class {bucket.__class__.__name__} has no attribute 'params' or 'params_list', please check!" |
| 146 | + ) | ||
| 147 | + logger.debug( | ||
| 148 | + f"megatron sync_grad_func enter: bucket_type={bucket.__class__.__name__}, " | ||
| 149 | + f"param_source={param_source}, bucket_param_count={len(param_attr)}, " | ||
| 150 | + f"target_param_count={len(monitor.param2name)}" | ||
| 151 | + ) | ||
| 144 | bucket_params_id_list = [id(params) for params in param_attr] | 152 | bucket_params_id_list = [id(params) for params in param_attr] |
| 145 | for param, name in monitor.param2name.items(): | 153 | for param, name in monitor.param2name.items(): |
| 146 | if id(param) not in bucket_params_id_list: | 154 | if id(param) not in bucket_params_id_list: |
| @@ -151,10 +159,13 @@ class OptimizerMon(object): | |||
| 151 | continue | 159 | continue |
| 152 | tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD) | 160 | tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD) |
| 153 | if tag is None: | 161 | if tag is None: |
| 162 | + logger.debug(f"megatron sync_grad_func skip param without tag: name={name}") | ||
| 154 | continue | 163 | continue |
| 155 | grad_dict[tag] = grad | 164 | grad_dict[tag] = grad |
| 156 | monitor.register_param_call_id("sync_grad_func", tag) | 165 | monitor.register_param_call_id("sync_grad_func", tag) |
| 166 | + logger.debug(f"megatron sync_grad_func collected grad_count={len(grad_dict)}") | ||
| 157 | get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre) | 167 | get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre) |
| 168 | + logger.debug(f"megatron sync_grad_func pre_metric_count={len(monitor.grad_context.pre)}") | ||
| 158 | out = sync_grad_func(bucket, *args, **kwargs) | 169 | out = sync_grad_func(bucket, *args, **kwargs) |
| 159 | return out | 170 | return out |
| 160 | 171 | ||
| @@ -162,6 +173,7 @@ class OptimizerMon(object): | |||
| 162 | 173 | ||
| 163 | try: | 174 | try: |
| 164 | from megatron.core.distributed.param_and_grad_buffer import Bucket | 175 | from megatron.core.distributed.param_and_grad_buffer import Bucket |
| 176 | + | ||
| 165 | self.origin_funcs.append(Bucket.start_grad_sync) | 177 | self.origin_funcs.append(Bucket.start_grad_sync) |
| 166 | self.bucket_class = Bucket | 178 | self.bucket_class = Bucket |
| 167 | Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync) | 179 | Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync) |
| @@ -172,6 +184,7 @@ class OptimizerMon(object): | |||
| 172 | 184 | ||
| 173 | try: | 185 | try: |
| 174 | from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup | 186 | from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup |
| 187 | + | ||
| 175 | self.origin_funcs.append(_ParamAndGradBucketGroup.start_grad_sync) | 188 | self.origin_funcs.append(_ParamAndGradBucketGroup.start_grad_sync) |
| 176 | self.bucket_class = _ParamAndGradBucketGroup | 189 | self.bucket_class = _ParamAndGradBucketGroup |
| 177 | _ParamAndGradBucketGroup.start_grad_sync = patch_sync(_ParamAndGradBucketGroup.start_grad_sync) | 190 | _ParamAndGradBucketGroup.start_grad_sync = patch_sync(_ParamAndGradBucketGroup.start_grad_sync) |
| @@ -207,7 +220,8 @@ class MegatronMixPrecisionOptimizerMon(OptimizerMon): | |||
| 207 | if not (hasattr(torch_opt, "float16_groups") and hasattr(torch_opt, "fp32_from_float16_groups")): | 220 | if not (hasattr(torch_opt, "float16_groups") and hasattr(torch_opt, "fp32_from_float16_groups")): |
| 208 | logger.warning( | 221 | logger.warning( |
| 209 | "megatron class Float16OptimizerWithFloat16Params should have float16_groups and " | 222 | "megatron class Float16OptimizerWithFloat16Params should have float16_groups and " |
| 210 | - "fp32_from_float16_groups, please check Megatron-LM version.") | 223 | + "fp32_from_float16_groups, please check Megatron-LM version." |
| 224 | + ) | ||
| 211 | else: | 225 | else: |
| 212 | for fp16_group, fp32_group in zip(torch_opt.float16_groups, torch_opt.fp32_from_float16_groups): | 226 | for fp16_group, fp32_group in zip(torch_opt.float16_groups, torch_opt.fp32_from_float16_groups): |
| 213 | for fp16_param, fp32_param in zip(fp16_group, fp32_group): | 227 | for fp16_param, fp32_param in zip(fp16_group, fp32_group): |
| @@ -219,10 +233,12 @@ class MegatronDistributedOptimizerMon(OptimizerMon): | |||
| 219 | if not (hasattr(torch_opt, "model_float16_groups") and hasattr(torch_opt, "shard_fp32_from_float16_groups")): | 233 | if not (hasattr(torch_opt, "model_float16_groups") and hasattr(torch_opt, "shard_fp32_from_float16_groups")): |
| 220 | logger.warning( | 234 | logger.warning( |
| 221 | "megatron class DistributedOptimizer should have model_float16_groups and " | 235 | "megatron class DistributedOptimizer should have model_float16_groups and " |
| 222 | - "shard_fp32_from_float16_groups, please check Megatron-LM version.") | 236 | + "shard_fp32_from_float16_groups, please check Megatron-LM version." |
| 237 | + ) | ||
| 223 | else: | 238 | else: |
| 224 | - for fp16_group, shard_fp32_group in zip(torch_opt.model_float16_groups, | 239 | + for fp16_group, shard_fp32_group in zip( |
| 225 | - torch_opt.shard_fp32_from_float16_groups): | 240 | + torch_opt.model_float16_groups, torch_opt.shard_fp32_from_float16_groups |
| 241 | + ): | ||
| 226 | for fp16_param, shard_fp32_param in zip(fp16_group, shard_fp32_group): | 242 | for fp16_param, shard_fp32_param in zip(fp16_group, shard_fp32_group): |
| 227 | self.fp16_to_fp32_param[fp16_param] = shard_fp32_param | 243 | self.fp16_to_fp32_param[fp16_param] = shard_fp32_param |
| 228 | 244 | ||
| @@ -373,14 +389,23 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon): | |||
| 373 | def patch_sync(reduce_func): | 389 | def patch_sync(reduce_func): |
| 374 | def wrapper(zero_optimizer, *args, **kwargs): | 390 | def wrapper(zero_optimizer, *args, **kwargs): |
| 375 | grad_dict = {} | 391 | grad_dict = {} |
| 392 | + logger.debug( | ||
| 393 | + f"deepspeed sync_grad_func enter: bucket_count={len(zero_optimizer.params_in_ipg_bucket)}, " | ||
| 394 | + f"target_param_count={len(monitor.param2name)}" | ||
| 395 | + ) | ||
| 376 | for i, param, _ in zero_optimizer.params_in_ipg_bucket: | 396 | for i, param, _ in zero_optimizer.params_in_ipg_bucket: |
| 377 | if isinstance(param, int): # for ds >= 0.17.0 | 397 | if isinstance(param, int): # for ds >= 0.17.0 |
| 378 | param = zero_optimizer.bit16_groups[i][param] | 398 | param = zero_optimizer.bit16_groups[i][param] |
| 379 | name = monitor.param2name[param] | 399 | name = monitor.param2name[param] |
| 380 | tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD) | 400 | tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD) |
| 401 | + if tag is None: | ||
| 402 | + logger.debug(f"deepspeed sync_grad_func skip param without tag: name={name}") | ||
| 403 | + continue | ||
| 381 | grad_dict[tag] = zero_optimizer.get_gradient_for_reduction(param) | 404 | grad_dict[tag] = zero_optimizer.get_gradient_for_reduction(param) |
| 382 | monitor.register_param_call_id("sync_grad_func", tag) | 405 | monitor.register_param_call_id("sync_grad_func", tag) |
| 406 | + logger.debug(f"deepspeed sync_grad_func collected grad_count={len(grad_dict)}") | ||
| 383 | get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre) | 407 | get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre) |
| 408 | + logger.debug(f"deepspeed sync_grad_func pre_metric_count={len(monitor.grad_context.pre)}") | ||
| 384 | out = reduce_func(zero_optimizer, *args, **kwargs) | 409 | out = reduce_func(zero_optimizer, *args, **kwargs) |
| 385 | return out | 410 | return out |
| 386 | 411 | ||
| @@ -388,16 +413,15 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon): | |||
| 388 | 413 | ||
| 389 | try: | 414 | try: |
| 390 | from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer | 415 | from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer |
| 391 | - self.origin_funcs = [ | 416 | + |
| 392 | - DeepSpeedZeroOptimizer.average_tensor, | 417 | + self.origin_funcs = [DeepSpeedZeroOptimizer.average_tensor, DeepSpeedZeroOptimizer.buffered_reduce_fallback] |
| 393 | - DeepSpeedZeroOptimizer.buffered_reduce_fallback | ||
| 394 | - ] | ||
| 395 | DeepSpeedZeroOptimizer.average_tensor = patch_sync(DeepSpeedZeroOptimizer.average_tensor) | 418 | DeepSpeedZeroOptimizer.average_tensor = patch_sync(DeepSpeedZeroOptimizer.average_tensor) |
| 396 | - DeepSpeedZeroOptimizer.buffered_reduce_fallback = \ | 419 | + DeepSpeedZeroOptimizer.buffered_reduce_fallback = patch_sync( |
| 397 | - patch_sync(DeepSpeedZeroOptimizer.buffered_reduce_fallback) | 420 | + DeepSpeedZeroOptimizer.buffered_reduce_fallback |
| 421 | + ) | ||
| 398 | monitor.enable_deepspeed = True | 422 | monitor.enable_deepspeed = True |
| 399 | logger.info('deepspeed enabled') | 423 | logger.info('deepspeed enabled') |
| 400 | - except Exception as e: | 424 | + except Exception: |
| 401 | monitor.enable_deepspeed = False | monitor.enable_deepspeed | 425 | monitor.enable_deepspeed = False | monitor.enable_deepspeed |
| 402 | logger.warning('Seems using deepspeed zero 1 or 2. But patch average tensor failed') | 426 | logger.warning('Seems using deepspeed zero 1 or 2. But patch average tensor failed') |
| 403 | 427 | ||
| @@ -406,6 +430,7 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon): | |||
| 406 | return | 430 | return |
| 407 | 431 | ||
| 408 | from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer | 432 | from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer |
| 433 | + | ||
| 409 | DeepSpeedZeroOptimizer.average_tensor = self.origin_funcs[0] | 434 | DeepSpeedZeroOptimizer.average_tensor = self.origin_funcs[0] |
| 410 | DeepSpeedZeroOptimizer.buffered_reduce_fallback = self.origin_funcs[1] | 435 | DeepSpeedZeroOptimizer.buffered_reduce_fallback = self.origin_funcs[1] |
| 411 | 436 | ||
| @@ -442,7 +467,7 @@ class OptimizerMonFactory: | |||
| 442 | "BF16_Optimizer": DeepSpeedZeroOptimizerStage0Mon, | 467 | "BF16_Optimizer": DeepSpeedZeroOptimizerStage0Mon, |
| 443 | "DeepSpeedZeroOptimizer": DeepSpeedZeroOptimizerStage1or2Mon, | 468 | "DeepSpeedZeroOptimizer": DeepSpeedZeroOptimizerStage1or2Mon, |
| 444 | "DeepSpeedZeroOptimizer_Stage3": DeepSpeedZeroOptimizerStage3Mon, | 469 | "DeepSpeedZeroOptimizer_Stage3": DeepSpeedZeroOptimizerStage3Mon, |
| 445 | - "Adam": OptimizerMon | 470 | + "Adam": OptimizerMon, |
| 446 | } | 471 | } |
| 447 | 472 | ||
| 448 | @staticmethod | 473 | @staticmethod |
| @@ -17,6 +17,7 @@ | |||
| 17 | import os.path | 17 | import os.path |
| 18 | import shutil | 18 | import shutil |
| 19 | import unittest | 19 | import unittest |
| 20 | +from collections import defaultdict | ||
| 20 | from unittest.mock import patch, MagicMock | 21 | from unittest.mock import patch, MagicMock |
| 21 | 22 | ||
| 22 | import pandas as pd | 23 | import pandas as pd |
| @@ -224,6 +225,71 @@ class TestTrainerMon(unittest.TestCase): | |||
| 224 | self.mon.dynamic_monitor(optimizer) | 225 | self.mon.dynamic_monitor(optimizer) |
| 225 | self.assertEqual(self.mon.config_timestamp, 123456) | 226 | self.assertEqual(self.mon.config_timestamp, 123456) |
| 226 | 227 | ||
| 228 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 229 | + def test_set_config_logs_normalized_summary(self, mock_debug): | ||
| 230 | + self.mon.set_config() | ||
| 231 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 232 | + self.assertTrue(any("set_config normalized" in msg for msg in logged_messages)) | ||
| 233 | + self.assertTrue(any("set_config writer setup" in msg for msg in logged_messages)) | ||
| 234 | + | ||
| 235 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 236 | + def test_generate_wgrad_metrics_logs_counts(self, mock_debug): | ||
| 237 | + self.mon.wg_distribution = True | ||
| 238 | + self.mon.ops = [] | ||
| 239 | + self.mon.eps = 1e-8 | ||
| 240 | + post_grad_dict = {"tag": torch.tensor([1.0])} | ||
| 241 | + | ||
| 242 | + with patch("msprobe.pytorch.monitor.module_hook.get_metrics") as mock_get_metrics: | ||
| 243 | + self.mon.generate_wgrad_metrics(post_grad_dict) | ||
| 244 | + | ||
| 245 | + mock_get_metrics.assert_called_once() | ||
| 246 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 247 | + self.assertTrue(any("generate_wgrad_metrics start" in msg for msg in logged_messages)) | ||
| 248 | + self.assertTrue(any("generate_wgrad_metrics finish" in msg for msg in logged_messages)) | ||
| 249 | + | ||
| 250 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 251 | + def test_write_grad_tb_logs_counts(self, mock_debug): | ||
| 252 | + self.mon.wg_distribution = True | ||
| 253 | + self.mon.monitor_mbs_grad = True | ||
| 254 | + self.mon.ops = [] | ||
| 255 | + self.mon.grad_context.pre = {"pre": torch.tensor([1.0])} | ||
| 256 | + self.mon.grad_context.post = {"post": torch.tensor([2.0])} | ||
| 257 | + self.mon.summary_writer = MagicMock() | ||
| 258 | + | ||
| 259 | + self.mon.write_grad_tb(step=3) | ||
| 260 | + | ||
| 261 | + self.mon.summary_writer.write_metrics.assert_any_call( | ||
| 262 | + self.mon.ops, self.mon.grad_context.pre, 3, 'grad_unreduced', use_micro_step=True | ||
| 263 | + ) | ||
| 264 | + self.mon.summary_writer.write_metrics.assert_any_call( | ||
| 265 | + self.mon.ops, self.mon.grad_context.post, 3, 'grad_reduced' | ||
| 266 | + ) | ||
| 267 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 268 | + self.assertTrue(any("write_grad_tb start" in msg for msg in logged_messages)) | ||
| 269 | + | ||
| 270 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 271 | + def test_register_chunk_logs_mapping(self, mock_debug): | ||
| 272 | + self.mon.config["targets"] = {"layer": {}} | ||
| 273 | + self.mon.param2name = defaultdict(str) | ||
| 274 | + self.mon.name2param = {} | ||
| 275 | + self.mon.origin2squash = {} | ||
| 276 | + self.mon.duplicate_param = {} | ||
| 277 | + self.mon.name2tag = {} | ||
| 278 | + self.mon.rank = 0 | ||
| 279 | + | ||
| 280 | + class DummyModule(torch.nn.Module): | ||
| 281 | + def __init__(self): | ||
| 282 | + super().__init__() | ||
| 283 | + self.layer = torch.nn.Linear(2, 2) | ||
| 284 | + | ||
| 285 | + module = DummyModule() | ||
| 286 | + self.mon._register_chunk(module, "0:") | ||
| 287 | + | ||
| 288 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 289 | + self.assertTrue(any("_register_chunk start" in msg for msg in logged_messages)) | ||
| 290 | + self.assertTrue(any("_register_chunk mapped param" in msg for msg in logged_messages)) | ||
| 291 | + self.assertTrue(any("_register_chunk finish" in msg for msg in logged_messages)) | ||
| 292 | + | ||
| 227 | @patch("msprobe.pytorch.monitor.module_hook.is_recomputation", return_value=False) | 293 | @patch("msprobe.pytorch.monitor.module_hook.is_recomputation", return_value=False) |
| 228 | @patch("msprobe.pytorch.monitor.module_hook.get_entropy_metric") | 294 | @patch("msprobe.pytorch.monitor.module_hook.get_entropy_metric") |
| 229 | @patch("msprobe.pytorch.monitor.module_hook.cal_qkt") | 295 | @patch("msprobe.pytorch.monitor.module_hook.cal_qkt") |
| @@ -459,9 +525,10 @@ class TestTrainerMon(unittest.TestCase): | |||
| 459 | mock_save_json.assert_not_called() | 525 | mock_save_json.assert_not_called() |
| 460 | mock_getmtime.assert_not_called() | 526 | mock_getmtime.assert_not_called() |
| 461 | 527 | ||
| 528 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 462 | @patch("msprobe.pytorch.monitor.module_hook.get_sign_matches") | 529 | @patch("msprobe.pytorch.monitor.module_hook.get_sign_matches") |
| 463 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") | 530 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") |
| 464 | - def test_hook_optimizer_mg_direction_branch(self, mock_get_metrics, mock_get_sign_matches): | 531 | + def test_hook_optimizer_mg_direction_branch(self, mock_get_metrics, mock_get_sign_matches, mock_debug): |
| 465 | class DummyOptimizer: | 532 | class DummyOptimizer: |
| 466 | def step(self): | 533 | def step(self): |
| 467 | return None | 534 | return None |
| @@ -507,6 +574,9 @@ class TestTrainerMon(unittest.TestCase): | |||
| 507 | pre_hook(optimizer, (), {}) | 574 | pre_hook(optimizer, (), {}) |
| 508 | mock_get_sign_matches.assert_called_once() | 575 | mock_get_sign_matches.assert_called_once() |
| 509 | self.assertTrue(torch.equal(context.param_mg_direction["p1"], mock_get_sign_matches.return_value)) | 576 | self.assertTrue(torch.equal(context.param_mg_direction["p1"], mock_get_sign_matches.return_value)) |
| 577 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 578 | + self.assertTrue(any("optimizer_pre_step_hook enter" in msg for msg in logged_messages)) | ||
| 579 | + self.assertTrue(any("optimizer_pre_step_hook metrics prepared" in msg for msg in logged_messages)) | ||
| 510 | 580 | ||
| 511 | def test_hook_optimizer_patches_instance_step_only(self): | 581 | def test_hook_optimizer_patches_instance_step_only(self): |
| 512 | class DummyOptimizer: | 582 | class DummyOptimizer: |
| @@ -528,9 +598,10 @@ class TestTrainerMon(unittest.TestCase): | |||
| 528 | result = optimizer1.step() | 598 | result = optimizer1.step() |
| 529 | self.assertEqual(result, "original_step") | 599 | self.assertEqual(result, "original_step") |
| 530 | 600 | ||
| 601 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 531 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") | 602 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") |
| 532 | @patch("torch.distributed.fsdp._runtime_utils._post_backward_hook") | 603 | @patch("torch.distributed.fsdp._runtime_utils._post_backward_hook") |
| 533 | - def test_patch_fsdp_post_backward_hook(self, mock_post_hook, mock_get_metrics): | 604 | + def test_patch_fsdp_post_backward_hook(self, mock_post_hook, mock_get_metrics, mock_debug): |
| 534 | from msprobe.core.common.const import MonitorConst as MC | 605 | from msprobe.core.common.const import MonitorConst as MC |
| 535 | 606 | ||
| 536 | class DummyFlatParam: | 607 | class DummyFlatParam: |
| @@ -573,10 +644,14 @@ class TestTrainerMon(unittest.TestCase): | |||
| 573 | args, _ = mock_get_metrics.call_args | 644 | args, _ = mock_get_metrics.call_args |
| 574 | grad_dict = args[1] | 645 | grad_dict = args[1] |
| 575 | self.assertIn("pre_grad_tag", grad_dict) | 646 | self.assertIn("pre_grad_tag", grad_dict) |
| 647 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 648 | + self.assertTrue(any("_post_backward_hook enter" in msg for msg in logged_messages)) | ||
| 649 | + self.assertTrue(any("_post_backward_hook collected pre_grad metrics" in msg for msg in logged_messages)) | ||
| 576 | 650 | ||
| 651 | + @patch("msprobe.pytorch.monitor.module_hook.logger.debug") | ||
| 577 | @patch("msprobe.pytorch.monitor.module_hook.importlib.reload") | 652 | @patch("msprobe.pytorch.monitor.module_hook.importlib.reload") |
| 578 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") | 653 | @patch("msprobe.pytorch.monitor.module_hook.get_metrics") |
| 579 | - def test_patch_fsdp2_foreach_reduce(self, mock_get_metrics, mock_reload): | 654 | + def test_patch_fsdp2_foreach_reduce(self, mock_get_metrics, mock_reload, mock_debug): |
| 580 | from msprobe.core.common.const import MonitorConst as MC | 655 | from msprobe.core.common.const import MonitorConst as MC |
| 581 | 656 | ||
| 582 | # 构造假的 fully_shard 子模块，放入 sys.modules，避免环境缺少该模块时报错 | 657 | # 构造假的 fully_shard 子模块，放入 sys.modules，避免环境缺少该模块时报错 |
| @@ -620,6 +695,9 @@ class TestTrainerMon(unittest.TestCase): | |||
| 620 | grad_dict = args[1] | 695 | grad_dict = args[1] |
| 621 | self.assertIn("pre_grad_tag", grad_dict) | 696 | self.assertIn("pre_grad_tag", grad_dict) |
| 622 | self.assertTrue(original_called["flag"]) | 697 | self.assertTrue(original_called["flag"]) |
| 698 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 699 | + self.assertTrue(any("foreach_reduce enter" in msg for msg in logged_messages)) | ||
| 700 | + self.assertTrue(any("foreach_reduce collected pre_grad metrics" in msg for msg in logged_messages)) | ||
| 623 | 701 | ||
| 624 | def test_is_recording_module(self): | 702 | def test_is_recording_module(self): |
| 625 | # 1. 初始化核心属性 | 703 | # 1. 初始化核心属性 |
| @@ -1,4 +1,5 @@ | |||
| 1 | import unittest | 1 | import unittest |
| 2 | +from collections import defaultdict | ||
| 2 | from types import SimpleNamespace | 3 | from types import SimpleNamespace |
| 3 | from unittest.mock import Mock, patch, MagicMock | 4 | from unittest.mock import Mock, patch, MagicMock |
| 4 | import sys | 5 | import sys |
| @@ -493,7 +494,8 @@ class TestPatchGradSyncWrapper(unittest.TestCase): | |||
| 493 | delattr(bucket, 'params_list') | 494 | delattr(bucket, 'params_list') |
| 494 | return bucket | 495 | return bucket |
| 495 | 496 | ||
| 496 | - def test_wrapper_passes_args_and_kwargs(self): | 497 | + @patch("msprobe.pytorch.monitor.optimizer_collect.logger.debug") |
| 498 | + def test_wrapper_passes_args_and_kwargs(self, mock_debug): | ||
| 497 | with patch.dict('sys.modules', { | 499 | with patch.dict('sys.modules', { |
| 498 | 'megatron': MagicMock(), | 500 | 'megatron': MagicMock(), |
| 499 | 'megatron.core': MagicMock(), | 501 | 'megatron.core': MagicMock(), |
| @@ -525,10 +527,55 @@ class TestPatchGradSyncWrapper(unittest.TestCase): | |||
| 525 | result = mock_bucket_module.Bucket.start_grad_sync( | 527 | result = mock_bucket_module.Bucket.start_grad_sync( |
| 526 | bucket, 'arg1', 'arg2', key1='kwarg1', key2='kwarg2' | 528 | bucket, 'arg1', 'arg2', key1='kwarg1', key2='kwarg2' |
| 527 | ) | 529 | ) |
| 528 | - | 530 | + |
| 529 | self.assertEqual(result, 'sync_result') | 531 | self.assertEqual(result, 'sync_result') |
| 530 | self.assertEqual(received_args, ['arg1', 'arg2']) | 532 | self.assertEqual(received_args, ['arg1', 'arg2']) |
| 531 | self.assertEqual(received_kwargs, [('key1', 'kwarg1'), ('key2', 'kwarg2')]) | 533 | self.assertEqual(received_kwargs, [('key1', 'kwarg1'), ('key2', 'kwarg2')]) |
| 534 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 535 | + self.assertTrue(any("megatron sync_grad_func enter" in msg for msg in logged_messages)) | ||
| 536 | + self.assertTrue(any("megatron sync_grad_func collected grad_count" in msg for msg in logged_messages)) | ||
| 537 | + | ||
| 538 | + | ||
| 539 | +class TestDeepSpeedPatchGradSync(unittest.TestCase): | ||
| 540 | + @patch("msprobe.pytorch.monitor.optimizer_collect.logger.debug") | ||
| 541 | + @patch("msprobe.pytorch.monitor.optimizer_collect.get_metrics") | ||
| 542 | + def test_patch_grad_sync_logs_bucket_info(self, mock_get_metrics, mock_debug): | ||
| 543 | + optimizer = DeepSpeedZeroOptimizerStage1or2Mon(MagicMock()) | ||
| 544 | + monitor = MagicMock() | ||
| 545 | + param = torch.nn.Parameter(torch.randn(2, 3)) | ||
| 546 | + monitor.param2name = defaultdict(str, {param: "test_param"}) | ||
| 547 | + monitor.name2tag = {"test_param": {MonitorConst.PRE_GRAD: "tag1"}} | ||
| 548 | + monitor.ops = [] | ||
| 549 | + monitor.eps = 1e-8 | ||
| 550 | + monitor.grad_context = SimpleNamespace(pre={}) | ||
| 551 | + monitor.register_param_call_id = MagicMock() | ||
| 552 | + | ||
| 553 | + with patch.dict('sys.modules', { | ||
| 554 | + 'deepspeed': MagicMock(), | ||
| 555 | + 'deepspeed.runtime': MagicMock(), | ||
| 556 | + 'deepspeed.runtime.zero': MagicMock(), | ||
| 557 | + 'deepspeed.runtime.zero.stage_1_and_2': MagicMock(), | ||
| 558 | + }): | ||
| 559 | + stage_module = sys.modules['deepspeed.runtime.zero.stage_1_and_2'] | ||
| 560 | + | ||
| 561 | + class DummyZeroOptimizer: | ||
| 562 | + average_tensor = staticmethod(lambda zero_optimizer, *args, **kwargs: "avg") | ||
| 563 | + buffered_reduce_fallback = staticmethod(lambda zero_optimizer, *args, **kwargs: "fallback") | ||
| 564 | + | ||
| 565 | + stage_module.DeepSpeedZeroOptimizer = DummyZeroOptimizer | ||
| 566 | + optimizer.patch_grad_sync(monitor) | ||
| 567 | + zero_optimizer = MagicMock() | ||
| 568 | + zero_optimizer.params_in_ipg_bucket = [(0, param, None)] | ||
| 569 | + zero_optimizer.bit16_groups = [[param]] | ||
| 570 | + zero_optimizer.get_gradient_for_reduction.return_value = torch.ones_like(param) | ||
| 571 | + | ||
| 572 | + result = stage_module.DeepSpeedZeroOptimizer.average_tensor(zero_optimizer) | ||
| 573 | + | ||
| 574 | + self.assertEqual(result, "avg") | ||
| 575 | + mock_get_metrics.assert_called_once() | ||
| 576 | + logged_messages = [call.args[0] for call in mock_debug.call_args_list] | ||
| 577 | + self.assertTrue(any("deepspeed sync_grad_func enter" in msg for msg in logged_messages)) | ||
| 578 | + self.assertTrue(any("deepspeed sync_grad_func collected grad_count" in msg for msg in logged_messages)) | ||
| 532 | 579 | ||
| 533 | 580 | ||
| 534 | class TestOptimizerMonFactory(unittest.TestCase): | 581 | class TestOptimizerMonFactory(unittest.TestCase): |
此条代码评论区间+648至+708
热路径上使用 f-string 格式化 debug 日志存在性能开销。optimizer_pre_step_hook 和 step_final_hook 在每个 optimizer step 都会被调用，其中包含多达 7 个 logger.debug() 调用，每个都使用 f-string 格式化。Python 的 f-string 在函数调用前就会被求值，即使日志级别高于 DEBUG 时也无法跳过格式化开销。优化器钩子中的循环监控每个 step，累积的字符串格式化开销在长时间训练中会比较显著。 建议对热路径上的 debug 日志使用惰性求值模式，或使用 % 格式化方式让 logging 框架延迟求值。对于频繁调用的钩子函数，也可以考虑将多个日志合并为一条以减少调用次数。
# 推荐方式：使用 %s 格式化，logging 框架内部按需求值 logger.debug("optimizer_pre_step_hook enter: step=%s, start_step=%s, step_interval=%s, optimizer_type=%s, wg_distribution=%s, mv_distribution=%s, ur_distribution=%s, mg_direction=%s", context.step, self.start_step, self.step_interval, optimizer.__class__.__name__, self.wg_distribution, self.mv_distribution, self.ur_distribution, self.mg_direction) # 或者将多条日志合并为一条 logger.debug( "optimizer_pre_step_hook: step=%s, grads=%s, mv_metrics=%s, param_metrics=%s", context.step, len(grad_dict) if self.wg_distribution else 'N/A', f'exp_avg={len(context.exp_avg_metric)}' if context.exp_avg_metric else 'N/A', f'pre={len(self.grad_context.pre)},post={len(self.grad_context.post)}' )