mindstudio-probe 8.1.2__py3-none-any.whl → 8.2.1__py3-none-any.whl

msprobe/pytorch/monitor/optimizer_collect.py

@@ -18,6 +18,7 @@ import torch
 
 from msprobe.pytorch.common.log import logger
 from msprobe.core.monitor.utils import MVResult
+from msprobe.pytorch.monitor.module_metric import get_metrics
 from msprobe.core.common.const import MonitorConst
 
 
@@ -26,6 +27,8 @@ class OptimizerMon(object):
         self.fp16_to_fp32_param = {}
         self.torch_opt = torch_opt
         self.state = {}
+        self.origin_funcs = []
+        self.bucket_class = None
 
     def narrow_from_flatten(self, param, flatten_state):
         return flatten_state
@@ -49,11 +52,13 @@ class OptimizerMon(object):
             if self.fp16_to_fp32_param and param not in self.fp16_to_fp32_param:
                 continue
             grad = param.main_grad if monitor.params_have_main_grad else param.grad
+            if grad.__class__.__name__ == 'DTensor':
+                grad = grad.to_local()
             element_in_cur_partition = self.fp16_to_fp32_param.get(param, param).numel()
             if param.numel() != element_in_cur_partition:
                 if first_param:
                     grad = grad.flatten()[-element_in_cur_partition:]
-                else: # supposed to be the last one
+                else:  # supposed to be the last one
                     grad = grad.flatten()[:element_in_cur_partition]
             first_param = False
 
@@ -120,6 +125,59 @@ class OptimizerMon(object):
                     monitor.ratio_heatmap_visualizer[name].pre_cal(ratio_dict[name])
         return MVResult(exp_avg=exp_avg_dict, exp_avg_sq=exp_avg_sq_dict, update=update_dict, ratio=ratio_dict)
     
+    def patch_grad_sync(self, monitor):
+        def patch_sync(sync_grad_func):
+            def wrapper(bucket):
+                grad_dict = {}
+                # Megatron between core_r0.6.0 and core_r0.8.0, this bucket is Bucket.
+                # When megatron is core_r0.9.0, this bucket is _ParamAndGradBucketGroup.
+                # In megatron version core_r0.9.0, func start_grad_sync from Bucket moved to _ParamAndGradBucketGroup.
+                bucket_params_id_list = [id(params) for params in bucket.params]
+                for param, name in monitor.param2name.items():
+                    if id(param) not in bucket_params_id_list:
+                        continue
+                    grad = param.main_grad if monitor.params_have_main_grad else param.grad
+                    if grad is None:
+                        logger.warning(f"grad is None: {name}, maybe something wrong happened.")
+                        continue
+                    tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD)
+                    if tag is None:
+                        continue
+                    grad_dict[tag] = grad
+                    monitor.register_param_call_id("sync_grad_func", tag)
+                get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre)
+                out = sync_grad_func(bucket)
+                return out
+
+            return wrapper
+    
+        try:
+            from megatron.core.distributed.param_and_grad_buffer import Bucket
+            self.origin_funcs.append(Bucket.start_grad_sync)
+            self.bucket_class = Bucket
+            Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync)
+            monitor.enable_megatron = True
+            logger.info("megatron version is >= core_r0.6.0 <= core_r0.8.0")
+        except ImportError:
+            monitor.enable_megatron = False
+
+        try:
+            from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup
+            self.origin_funcs.append(_ParamAndGradBucketGroup.start_grad_sync)
+            self.bucket_class = _ParamAndGradBucketGroup
+            _ParamAndGradBucketGroup.start_grad_sync = patch_sync(_ParamAndGradBucketGroup.start_grad_sync)
+            monitor.enable_megatron = True
+            logger.info("megatron version is > core_r0.8.0 <= core_r0.9.0")
+        except ImportError:
+            monitor.enable_megatron = False | monitor.enable_megatron
+
+    def restore_grad_sync(self, monitor):
+        if not monitor.enable_megatron:
+            return
+        
+        self.bucket_class.start_grad_sync = self.origin_funcs[0]
+        
+
     def _get_single_state(self, torch_opt):
         state = {}
         if hasattr(torch_opt, 'param_to_cpu_states_map'):
@@ -131,7 +189,7 @@ class OptimizerMon(object):
         self.state.update(state)
 
 
-class MixPrecisionOptimizerMon(OptimizerMon):
+class MegatronMixPrecisionOptimizerMon(OptimizerMon):
     """
     混合精度优化器监控类。在混合精度训练中监控和管理优化器。
     混合精度训练通过适当降低某些计算的精度来加速训练过程并减少内存消耗。
@@ -161,7 +219,7 @@ class MegatronChainedDistributedOptimizerMon(MegatronDistributedOptimizerMon):
             super().map_fp16_to_fp32_param(opt)
 
 
-class MegatronChainedMixPrecisionOptimizerMon(MixPrecisionOptimizerMon):
+class MegatronChainedMixPrecisionOptimizerMon(MegatronMixPrecisionOptimizerMon):
     def map_fp16_to_fp32_param(self, torch_opt):
         for opt in torch_opt.chained_optimizers:
             super().map_fp16_to_fp32_param(opt)
@@ -248,6 +306,12 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
             grad_dict[tag] = grad
 
         return grad_dict
+    
+    def patch_grad_sync(self, monitor):
+        pass
+
+    def restore_grad_sync(self, monitor):
+        pass
 
 
 class DeepSpeedZeroOptimizerStage0Mon(DeepSpeedZeroOptimizerMon):
@@ -291,6 +355,47 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon):
                     break
 
 
+    def patch_grad_sync(self, monitor):
+        def patch_sync(reduce_func):
+            def wrapper(zero_optimizer, *args, **kwargs):
+                grad_dict = {}
+                for i, param, _ in zero_optimizer.params_in_ipg_bucket:
+                    if isinstance(param, int): # for ds >= 0.17.0
+                        param = zero_optimizer.bit16_groups[i][param]
+                    name = monitor.param2name[param]
+                    tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD)
+                    grad_dict[tag] = zero_optimizer.get_gradient_for_reduction(param)
+                    monitor.register_param_call_id("sync_grad_func", tag)
+                get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre)
+                out = reduce_func(zero_optimizer, *args, **kwargs)
+                return out
+            
+            return wrapper
+        try:
+            from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
+            self.origin_funcs = [
+                DeepSpeedZeroOptimizer.average_tensor, 
+                DeepSpeedZeroOptimizer.buffered_reduce_fallback
+                ]
+            DeepSpeedZeroOptimizer.average_tensor = patch_sync(DeepSpeedZeroOptimizer.average_tensor)
+            DeepSpeedZeroOptimizer.buffered_reduce_fallback = \
+                patch_sync(DeepSpeedZeroOptimizer.buffered_reduce_fallback)
+            monitor.enable_deepspeed = True
+            logger.info('deepspeed enabled')
+        except Exception as e:
+            monitor.enable_deepspeed = False | monitor.enable_deepspeed
+            logger.warning('Seems using deepspeed zero 1 or 2. But patch average tensor failed')
+
+    def restore_grad_sync(self, monitor):
+        if not monitor.enable_deepspeed:
+            return
+    
+        from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
+        DeepSpeedZeroOptimizer.average_tensor = self.origin_funcs[0]
+        DeepSpeedZeroOptimizer.buffered_reduce_fallback = self.origin_funcs[1]
+
+
+
 class DeepSpeedZeroOptimizerStage3Mon(DeepSpeedZeroOptimizerMon):
     def __init__(self, torch_opt):
         super().__init__(torch_opt)
@@ -314,7 +419,7 @@ class DeepSpeedZeroOptimizerStage3Mon(DeepSpeedZeroOptimizerMon):
 class OptimizerMonFactory:
     _optimizer_mon_map = {
         "FP32Optimizer": OptimizerMon,
-        "Float16OptimizerWithFloat16Params": MixPrecisionOptimizerMon,
+        "Float16OptimizerWithFloat16Params": MegatronMixPrecisionOptimizerMon,
         "DistributedOptimizer": MegatronDistributedOptimizerMon,
         "SwapDistributedOptimizer": MegatronDistributedOptimizerMon,
         "ChainedDistributedOptimizer": MegatronChainedDistributedOptimizerMon,

msprobe/pytorch/online_dispatch/dispatch.py

@@ -17,7 +17,7 @@ import json
 import os
 import time
 import multiprocessing
-from multiprocessing import Pool
+from multiprocessing import Pool, Lock
 
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -39,6 +39,7 @@ from msprobe.pytorch.online_dispatch.utils import get_callstack, data_to_cpu, ge
 from msprobe.pytorch.online_dispatch.compare import Comparator
 from msprobe.core.common.utils import check_str_param, safe_get_value
 
+child_global_lock = None
 current_time = time.strftime("%Y%m%d%H%M%S")
 RESULT_FILE_NAME = "accuracy_checking_result_" + current_time + ".csv"
 DETAILS_FILE_NAME = "accuracy_checking_details_" + current_time + ".csv"
@@ -86,14 +87,14 @@ class PtdbgDispatch(TorchDispatchMode):
         yaml_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "torch_ops_config.yaml")
         self.get_ops(yaml_path)
 
-        self.lock = None
+        self.lock = Lock() if process_num > 0 else None
         max_process_num = max(int((multiprocessing.cpu_count() + 1) // Const.CPU_QUARTER), 1)
         if process_num > max_process_num:
             logger.error(f"process_num should be less than or equal to {max_process_num}, but got {process_num}!")
             raise DispatchException(f'process_num should be less than or equal to {max_process_num}, '
                                     f'but got {process_num}!')
         if process_num > 0:
-            self.pool = Pool(process_num)
+            self.pool = Pool(process_num, initializer=self._init_child_process, initargs=(self.lock,))
         if debug:
             logger.info(f'Main pid:{os.getpid()} device:{self.device_id} dump_list:{self.dump_api_list} '
                         f'dump_mode:{self.dump_mode} cpu_path[{self.root_cpu_path}], npu_path[{self.root_npu_path}], '
@@ -114,18 +115,17 @@ class PtdbgDispatch(TorchDispatchMode):
                 logger.error("Please check train log, An exception may have occurred!")
                 return
             check_file_or_directory_path(summary_path, False)
-            fp_handle = FileOpen(summary_path, "r")
-            while True:
-                json_line_data = fp_handle.readline()
-                if json_line_data == '\n':
-                    continue
-                if len(json_line_data) == 0:
-                    break
-                msg = json.loads(json_line_data)
-                if len(msg) < 2:
-                    raise ValueError("JSON data does not contain enough elements. Expected at least 2 elements.")
-                self.all_summary[msg[0]] = msg[1]
-            fp_handle.close()
+            with FileOpen(summary_path, "r") as fp_handle:
+                while True:
+                    json_line_data = fp_handle.readline()
+                    if json_line_data == '\n':
+                        continue
+                    if len(json_line_data) == 0:
+                        break
+                    msg = json.loads(json_line_data)
+                    if len(msg) < 2:
+                        raise ValueError("JSON data does not contain enough elements. Expected at least 2 elements.")
+                    self.all_summary[msg[0]] = msg[1]
 
         if self.debug_flag:
             input_num = 0
@@ -163,11 +163,16 @@ class PtdbgDispatch(TorchDispatchMode):
 
         call_stack = get_callstack()
         self.call_stack_list.append(call_stack)
-        self.api_index += 1
-        if aten_api not in self.single_api_index_dict:
-            self.single_api_index_dict[aten_api] = 1
-        else:
-            self.single_api_index_dict[aten_api] += 1
+
+        self.lock.acquire() if self.process_num > 0 else None
+        try:
+            self.api_index += 1
+            if aten_api not in self.single_api_index_dict:
+                self.single_api_index_dict[aten_api] = 1
+            else:
+                self.single_api_index_dict[aten_api] += 1
+        finally:
+            self.lock.release() if self.process_num > 0 else None
 
         run_param = self.get_run_param(aten_api, func.__name__, aten_api_overload_name)
 
@@ -180,7 +185,7 @@ class PtdbgDispatch(TorchDispatchMode):
         cpu_kwargs = []
         data_to_cpu(args, 0, cpu_args)
         data_to_cpu(kwargs, 0, cpu_kwargs)
-        
+
         cpu_args = safe_get_value(cpu_args, 0, "cpu_args")
         cpu_kwargs = safe_get_value(cpu_kwargs, 0, "cpu_kwargs")
 
@@ -194,7 +199,12 @@ class PtdbgDispatch(TorchDispatchMode):
             try:
                 cpu_out = func(*cpu_args, **cpu_kwargs)
             except RuntimeError as e:
-                self.api_index -= 1
+                self.lock.acquire() if self.process_num > 0 else None
+                try:
+                    self.api_index -= 1
+                    self.single_api_index_dict[aten_api] -= 1
+                finally:
+                    self.lock.release() if self.process_num > 0 else None
                 logger.warning(f"RuntimeError: {e}")
                 logger.warning(f"This aten_api {aten_api} does not support running on cpu, so skip it.")
                 return npu_out
@@ -215,7 +225,7 @@ class PtdbgDispatch(TorchDispatchMode):
             run_param.process_flag = True
             if self.check_fun(func, run_param):
                 data_info = DisPatchDataInfo(cpu_args, cpu_kwargs, self.all_summary, None, npu_out_cpu, cpu_out,
-                                             self.lock)
+                                             child_global_lock)
                 self.pool.apply_async(func=dispatch_multiprocess, args=(run_param, data_info),
                                       error_callback=error_call)
             else:
@@ -233,12 +243,20 @@ class PtdbgDispatch(TorchDispatchMode):
                     return True
         return False
 
+    @staticmethod
+    def _init_child_process(lock):
+        global child_global_lock
+        child_global_lock = lock
+
     def get_dir_name(self, tag):
         # guarantee file uniqueness
         time.sleep(1)
-        time_now = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+        # 时间格式：年-月-日-时-分-秒-毫秒（精确到千分之一秒）
+        time_now = time.strftime("%Y%m%d%H%M%S%f", time.localtime(time.time()))[:-3]  # 取前3位毫秒
+
         if tag is None or not isinstance(tag, str):
             logger.warning('There is not tag or the type of tag is not string.')
+            # 目录名格式：msprobe_rank{设备ID}_{毫秒时间戳}
             dir_name = f'msprobe_rank{self.device_id}_{time_now}'
         else:
             dir_name = f'msprobe_{tag}_rank{self.device_id}_{time_now}'

msprobe/pytorch/online_dispatch/dump_compare.py

@@ -21,7 +21,7 @@ from datetime import datetime, timezone
 import torch
 from msprobe.core.common.const import Const
 from msprobe.core.common.decorator import recursion_depth_decorator
-from msprobe.core.common.file_utils import FileOpen, save_npy, save_json, check_link, remove_path
+from msprobe.core.common.file_utils import FileOpen, save_npy, save_json, remove_path, check_link
 from msprobe.pytorch.common.log import logger
 
 

msprobe/pytorch/parse_tool/lib/visualization.py

@@ -83,4 +83,3 @@ class Visualization:
                         self.util.log.info("\nStatistic Info:")
                         title_printed = True
                     self.util.log.info(summery_info)
-            pkl_handle.close()

msprobe/pytorch/pt_config.py

@@ -35,48 +35,15 @@ from msprobe.pytorch.hook_module.utils import get_ops
 class TensorConfig(BaseConfig):
     def __init__(self, json_config):
         super().__init__(json_config)
-        self.online_run_ut = json_config.get("online_run_ut", False)
-        self.nfs_path = json_config.get("nfs_path", "")
-        self.host = json_config.get("host", "")
-        self.port = json_config.get("port", -1)
-        self.tls_path = json_config.get("tls_path", "./")
-        self.online_run_ut_recompute = json_config.get("online_run_ut_recompute", False)
         self.check_config()
         self._check_summary_mode()
         self._check_file_format()
-        if self.online_run_ut:
-            self._check_online_run_ut()
+
 
     def _check_file_format(self):
         if self.file_format is not None and self.file_format not in ["npy", "bin"]:
             raise Exception("file_format is invalid")
 
-    def _check_online_run_ut(self):
-        if not isinstance(self.online_run_ut, bool):
-            raise Exception(f"online_run_ut: {self.online_run_ut} is invalid.")
-
-        if not isinstance(self.online_run_ut_recompute, bool):
-            raise Exception(f"online_run_ut_recompute: {self.online_run_ut_recompute} is invalid.")
-
-        if self.nfs_path:
-            check_file_or_directory_path(self.nfs_path, isdir=True)
-            return
-
-        if self.tls_path:
-            check_file_or_directory_path(self.tls_path, isdir=True)
-            check_file_or_directory_path(os.path.join(self.tls_path, "client.key"))
-            check_file_or_directory_path(os.path.join(self.tls_path, "client.crt"))
-            check_file_or_directory_path(os.path.join(self.tls_path, "ca.crt"))
-            crl_path = os.path.join(self.tls_path, "crl.pem")
-            if os.path.exists(crl_path):
-                check_file_or_directory_path(crl_path)
-
-        if not isinstance(self.host, str) or not re.match(Const.ipv4_pattern, self.host):
-            raise Exception(f"host: {self.host} is invalid.")
-
-        if not isinstance(self.port, int) or not (0 < self.port <= 65535):
-            raise Exception(f"port: {self.port} is invalid, port range 0-65535.")
-
 
 class StatisticsConfig(BaseConfig):
     def __init__(self, json_config):
@@ -251,12 +218,7 @@ class RunUTConfig(BaseConfig):
         self.white_list = json_config.get("white_list", Const.DEFAULT_LIST)
         self.black_list = json_config.get("black_list", Const.DEFAULT_LIST)
         self.error_data_path = json_config.get("error_data_path", Const.DEFAULT_PATH)
-        self.is_online = json_config.get("is_online", False)
-        self.nfs_path = json_config.get("nfs_path", "")
-        self.host = json_config.get("host", "")
-        self.port = json_config.get("port", -1)
-        self.rank_list = json_config.get("rank_list", Const.DEFAULT_LIST)
-        self.tls_path = json_config.get("tls_path", "./")
+
         self.check_run_ut_config()
 
     @classmethod
@@ -274,22 +236,11 @@ class RunUTConfig(BaseConfig):
         if not os.path.exists(error_data_path):
             raise Exception("error_data_path: %s does not exist" % error_data_path)
 
-    @classmethod
-    def check_nfs_path_config(cls, nfs_path):
-        if nfs_path:
-            FileChecker(nfs_path, FileCheckConst.DIR, FileCheckConst.READ_ABLE).common_check()
-
-    @classmethod
-    def check_tls_path_config(cls, tls_path):
-        if tls_path:
-            FileChecker(tls_path, FileCheckConst.DIR, FileCheckConst.READ_ABLE).common_check()
 
     def check_run_ut_config(self):
         RunUTConfig.check_filter_list_config(Const.WHITE_LIST, self.white_list)
         RunUTConfig.check_filter_list_config(Const.BLACK_LIST, self.black_list)
         RunUTConfig.check_error_data_path_config(self.error_data_path)
-        RunUTConfig.check_nfs_path_config(self.nfs_path)
-        RunUTConfig.check_tls_path_config(self.tls_path)
 
 
 class GradToolConfig(BaseConfig):

msprobe/pytorch/pytorch_service.py

@@ -15,18 +15,14 @@
 
 from msprobe.core.common.utils import Const
 from msprobe.core.service import BaseService
-from msprobe.pytorch.attl_manager import ATTLManager
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import get_rank_if_initialized, torch_version_above_or_equal_2
+from msprobe.pytorch.common.utils import get_rank_if_initialized
 from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
-from msprobe.pytorch.hook_module.api_register import get_api_register, ApiTemplate
+from msprobe.pytorch.hook_module.api_register import get_api_register, ApiTemplate, redirect_wait
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
-from msprobe.pytorch.hook_module.jit_script_wrapper import wrap_jit_script_func
 from msprobe.pytorch.hook_module.pt_hook_manager import PytorchHookManager
 from msprobe.pytorch.hook_module.register_optimizer_hook import register_optimizer_hook
-
-if torch_version_above_or_equal_2:
-    from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.dump_dispatch import run_ut_dispatch
+from msprobe.pytorch.hook_module.script_wrapper import wrap_script_func, preprocess_func
 
 
 class PytorchService(BaseService):
@@ -45,27 +41,24 @@ class PytorchService(BaseService):
         self.logger = logger
         self.api_register = get_api_register()
         self.module_processor = ModuleProcesser(self.data_collector.scope)
-        self.attl_manager = ATTLManager(self.config)
-        self.hook_manager = PytorchHookManager(self.data_collector, self.config, self.attl_manager)
+        self.hook_manager = PytorchHookManager(self.data_collector, self.config)
         self.api_template = ApiTemplate
 
     def _register_hook(self):
-        self.attl_manager.attl_init()
         if self._is_mix_level:
             register_optimizer_hook(self.data_collector)
 
     def _register_api_hook(self):
+        preprocess_func()
         super()._register_api_hook()
-        wrap_jit_script_func()
+        wrap_script_func()
+        redirect_wait()
 
     def _register_module_hook(self):
         ModuleProcesser.enable_module_dump = True
         self.module_processor.register_module_hook(self.model, self.build_hook)
         self.logger.info(f"The module {self.config.task} hook function is successfully mounted to the model.")
 
-    def _run_ut_dispatch(self, status):
-        if torch_version_above_or_equal_2:
-            run_ut_dispatch(self.attl_manager.attl, status, self.config.online_run_ut_recompute)
 
     def _reset_status(self):
         super()._reset_status()