mindstudio-probe 1.1.1__py3-none-any.whl → 1.2.2__py3-none-any.whl

msprobe/pytorch/monitor/distributed/wrap_distributed.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -13,16 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import os
 import re
-import inspect
 
 import torch
-import torch.nn as nn
 import torch.distributed as dist
+import torch.nn as nn
 
-from msprobe.core.common.file_utils import load_yaml
 from msprobe.core.common.const import MonitorConst
+from msprobe.core.common.file_utils import load_yaml
 from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_writer_tag_name
 
 try:
@@ -57,10 +57,13 @@ class DistributedOPTemplate(nn.Module):
         super(DistributedOPTemplate, self).__init__()
         self.op_name_ = str(op_name)
         self.__name__ = self.op_name_
+        self.cc_hooks = []
         for pre_hook in pre_hooks:
-            self.register_forward_pre_hook(pre_hook, with_kwargs=True)
+            handle = self.register_forward_pre_hook(pre_hook, with_kwargs=True)
+            self.cc_hooks.append(handle)
         for hook in post_hooks:
-            self.register_forward_hook(hook, with_kwargs=True)
+            handle = self.register_forward_hook(hook, with_kwargs=True)
+            self.cc_hooks.append(handle)
 
     def forward(self, *args, **kwargs):
         return distributed_func.get(self.op_name_)(*args, **kwargs)
@@ -120,8 +123,11 @@ class ApiRegistry:
 
     def initialize_hook(self, pre_hooks, post_hooks):
         self.store_ori_attr(dist, get_distributed_ops(), self.distributed_attr_origin)
+        cc_hooks = []
         for op_name in get_distributed_ops():
             self.distributed_attr_hooked[op_name] = DistributedOPTemplate(op_name, pre_hooks, post_hooks)
+            cc_hooks.extend(self.distributed_attr_hooked[op_name].cc_hooks)
+        return cc_hooks
 
 
 def get_process_group(process_group):
@@ -237,15 +243,20 @@ def create_hooks(context, monitor):
         args = args + tuple(kwargs.values())
         if out:  # async
             if isinstance(out, dist.Work):
-                PENDING_ASYNC_CC_BY_HANDLE[out] = create_async_callback_func(context[module.op_name_],
-                                                                             module.op_name_,
-                                                                             monitor.ops, args, MonitorConst.PREFIX_POST)
-            elif isinstance(out, list): # batch_isend_irecv
+                PENDING_ASYNC_CC_BY_HANDLE[out] = create_async_callback_func(
+                    context[module.op_name_],
+                    module.op_name_,
+                    monitor.ops, args,
+                    MonitorConst.PREFIX_POST
+                )
+            elif isinstance(out, list):  # batch_isend_irecv
                 for out_element in out:
-                    PENDING_ASYNC_CC_BY_HANDLE[out_element] = create_async_callback_func(context[module.op_name_],
-                                                                                         module.op_name_,
-                                                                                         monitor.ops, args,
-                                                                                         MonitorConst.PREFIX_POST)
+                    PENDING_ASYNC_CC_BY_HANDLE[out_element] = create_async_callback_func(
+                        context[module.op_name_],
+                        module.op_name_,
+                        monitor.ops, args,
+                        MonitorConst.PREFIX_POST
+                    )
             return out
         catch_data(context[module.op_name_], module.op_name_, monitor.ops, args, MonitorConst.PREFIX_POST)
         return out

msprobe/pytorch/monitor/features.py

@@ -15,7 +15,7 @@
 
 import torch
 from torch.autograd.functional import jacobian
-from msprobe.core.common.log import logger
+from msprobe.pytorch.common.log import logger
 
 
 @torch.no_grad()
@@ -30,12 +30,12 @@ def get_min(x: torch.tensor):
 
 @torch.no_grad()
 def get_mean(x: torch.tensor):
-    return torch.mean(x)
+    return torch.mean(x.to(torch.float64))
 
 
 @torch.no_grad()
 def get_norm(x: torch.tensor):
-    return torch.norm(x, p=2)
+    return torch.norm(x.to(torch.float64), p=2)
 
 
 @torch.no_grad()