mindstudio-probe 8.3.3__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/core/monitor_v2/cc.py

@@ -0,0 +1,287 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Any, Dict, List, Tuple
+
+from msprobe.core.common.log import logger
+from msprobe.core.monitor_v2.base import BaseMonitorV2
+
+
+class _PTCCContextV2:
+    """
+    Lightweight communication context for PyTorch cc_distribution monitor_v2.
+
+    结构与 pytorch.monitor.module_hook.CommunicationContext 一致：
+      - data: {tag: {op: aggregated_tensor_or_scalar}}
+      - aggregate(): 对同一 tag/op 下的多次调用做聚合（min/max/mean/...）
+    """
+
+    def __init__(self) -> None:
+        self.data: Dict[str, Dict[str, Any]] = {}
+
+    @staticmethod
+    def _agg(data: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
+        from msprobe.pytorch.monitor.distributed.wrap_distributed import op_aggregate
+
+        aggregated_data: Dict[str, Dict[str, Any]] = {}
+        for tag, op2tensorlist in data.items():
+            aggregated_data[tag] = {}
+            for op, tensorlist in op2tensorlist.items():
+                aggregated_data[tag][op] = op_aggregate(op, tensorlist)
+        return aggregated_data
+
+    def reset(self) -> None:
+        self.data = {}
+
+    def aggregate(self) -> None:
+        self.data = self._agg(self.data)
+
+
+class _MSCCContextV2:
+    """
+    MindSpore 版本的通信 context，结构与 _PTCCContextV2 相同，
+    但使用 mindspore.wrap_distributed.op_aggregate。
+    """
+
+    def __init__(self) -> None:
+        self.data: Dict[str, Dict[str, Any]] = {}
+
+    @staticmethod
+    def _agg(data: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
+        from msprobe.mindspore.monitor.distributed.wrap_distributed import op_aggregate
+
+        aggregated_data: Dict[str, Dict[str, Any]] = {}
+        for tag, op2tensorlist in data.items():
+            aggregated_data[tag] = {}
+            for op, tensorlist in op2tensorlist.items():
+                aggregated_data[tag][op] = op_aggregate(op, tensorlist)
+        return aggregated_data
+
+    def reset(self) -> None:
+        self.data = {}
+
+    def aggregate(self) -> None:
+        self.data = self._agg(self.data)
+
+
+class _BaseCCMonitorV2(BaseMonitorV2, ABC):
+    """
+    通信算子监控（cc_distribution）的通用 v2 基类。
+
+    框架相关的细节（环境检查 / hook 注册 / API 恢复等）
+    由子类实现，避免在子类中再做 FmkAdp.fmk 等框架判断。
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._patched: bool = False
+        self._cc_handles: List[Any] = []
+        self.cc_context: Dict[str, Any] = defaultdict(self._create_context)
+
+        # cc 配置
+        self.cc_codeline: List[str] = []
+        self.cc_log_only: bool = False
+        self.cc_pre_hook: bool = False
+        self.cc_logged_stack: Dict[str, set] = defaultdict(set)
+        # 保持与 TrainerMon 同名，供 create_hooks 使用
+        self.module_rank_list: List[int] = []
+
+    @property
+    def ops(self) -> List[str]:
+        return self._ops
+
+    def set_config(self, config: Dict[str, Any]) -> None:
+        """
+        配置示例（monitors.cc）：
+          {
+            "enabled": true,
+            "ops": ["min", "max", "mean", "norm", "nans"],
+            "eps": 1e-8,
+            "cc_codeline": [],
+            "cc_log_only": false,
+            "cc_pre_hook": false,
+            "module_ranks": []
+          }
+        """
+        super().set_config(config)
+        self.config = config
+        # 兼容 cc_distribution 结构：如果嵌套提供，则取内部 dict
+        cc_cfg = config.get("cc_distribution", config)
+        self.cc_codeline = cc_cfg.get("cc_codeline", [])
+        self.cc_log_only = bool(cc_cfg.get("cc_log_only", False))
+        self.cc_pre_hook = bool(cc_cfg.get("cc_pre_hook", False))
+        self.module_rank_list = list(cc_cfg.get("module_ranks", []))
+
+    def start(self, model: Any = None, optimizer: Any = None, **context: Any) -> None:
+        self.set_context(**context)
+        if self._patched:
+            return
+
+        if not self._is_env_ready():
+            return
+
+        api_register, create_hooks = self._get_wrap_module()
+        pre_hooks, post_hooks = create_hooks(context=self.cc_context, monitor=self)
+        self._cc_handles = api_register.initialize_hook(pre_hooks, post_hooks)
+        api_register.redirect_api()
+        self._patched = True
+
+    def stop(self) -> None:
+        if not self._patched:
+            return
+        try:
+            self._restore_api()
+        except Exception as exc:
+            logger.warning(f"[monitor_v2] Failed to restore communication APIs: {exc}")
+        self._patched = False
+        self._cc_handles = []
+        for _, ctx in self.cc_context.items():
+            ctx.reset()
+        self.cc_context.clear()
+        self._rows = []
+
+    def collect(self) -> Dict[str, Any] | None:
+        """
+        聚合各通信算子的监控数据，并转为 rows。
+        """
+        if not self._patched or not self.cc_context:
+            return None
+
+        metric_dict: Dict[str, Dict[str, Any]] = {}
+        for ctx in self.cc_context.values():
+            ctx.aggregate()
+            # TrainerMon 中的行为一致：后写入的 op/tag 覆盖前面的同名项
+            for tag, op2val in ctx.data.items():
+                metric_dict.setdefault(tag, {}).update(op2val)
+            ctx.reset()
+
+        self._after_collect(metric_dict)
+
+        if not metric_dict:
+            return None
+
+        self._rows = []
+        for tag, op2val in metric_dict.items():
+            row = {
+                "module_name": tag,
+                "scope": "comm",
+                "stats": op2val,
+            }
+            self._rows.append(row)
+
+        return super().collect()
+
+    # ---- 子类需要实现的抽象方法 ----
+
+    @abstractmethod
+    def _create_context(self) -> Any:
+        ...
+
+    @abstractmethod
+    def _is_env_ready(self) -> bool:
+        ...
+
+    @abstractmethod
+    def _get_wrap_module(self) -> Tuple[Any, Any]:
+        """
+        返回 (api_register, create_hooks) 两个对象，供 start() 使用。
+        """
+
+    @abstractmethod
+    def _restore_api(self) -> None:
+        ...
+
+    def _after_collect(self, metric_dict: Dict[str, Dict[str, Any]]) -> None:
+        """
+        子类可在 collect 之后做一些额外操作（如 MindSpore 重置 idx）。
+        """
+        return
+
+
+class PyTorchCCMonitorV2(_BaseCCMonitorV2):
+    """
+    PyTorch 通信算子监控（cc_distribution）的 v2 实现。
+    """
+
+    def _create_context(self) -> Any:
+        return _PTCCContextV2()
+
+    def _is_env_ready(self) -> bool:
+        # 仅在分布式环境下挂通信 hook
+        try:
+            import torch.distributed as dist  # type: ignore[import]
+        except Exception:
+            return False
+        return dist.is_initialized()
+
+    def _get_wrap_module(self) -> Tuple[Any, Any]:
+        from msprobe.pytorch.monitor.distributed.wrap_distributed import (
+            api_register,
+            create_hooks,
+        )
+
+        return api_register, create_hooks
+
+    def _restore_api(self) -> None:
+        from msprobe.pytorch.monitor.distributed.wrap_distributed import api_register
+
+        api_register.restore_api()
+
+
+class MindSporeCCMonitorV2(_BaseCCMonitorV2):
+    """
+    MindSpore 通信算子监控（cc_distribution）的 v2 实现。
+
+    骨架对齐 mindspore/monitor.module_hook 中 cc_distribution 的逻辑。
+    """
+
+    def _create_context(self) -> Any:
+        return _MSCCContextV2()
+
+    def _is_env_ready(self) -> bool:
+        try:
+            from mindspore import communication  # type: ignore[import]
+        except Exception:
+            return False
+        if not getattr(communication, "GlobalComm", None) or not communication.GlobalComm.INITED:
+            return False
+        return True
+
+    def _get_wrap_module(self) -> Tuple[Any, Any]:
+        from msprobe.mindspore.monitor.distributed.wrap_distributed import (
+            api_register,
+            create_hooks,
+        )
+
+        return api_register, create_hooks
+
+    def _restore_api(self) -> None:
+        from msprobe.mindspore.monitor.distributed.wrap_distributed import api_register
+
+        api_register.restore_api()
+
+    def _after_collect(self, metric_dict: Dict[str, Dict[str, Any]]) -> None:
+        # reset communication op call index for next step
+        try:
+            from msprobe.mindspore.monitor.distributed.wrap_distributed import api_register
+
+            api_register.reset_idx()
+        except Exception as exc:
+            logger.warning(f"[monitor_v2] Failed to reset MindSpore communication index: {exc}")

msprobe/core/monitor_v2/factory.py

@@ -0,0 +1,81 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Type
+
+from msprobe.core.common.log import logger
+from msprobe.core.monitor_v2.module import (
+    MindSporeModuleMonitorV2,
+    PyTorchModuleMonitorV2,
+)
+from msprobe.core.monitor_v2.cc import PyTorchCCMonitorV2, MindSporeCCMonitorV2
+from msprobe.core.monitor_v2.optimizer import (
+    MindSporeOptimizerMonitorV2,
+    PyTorchOptimizerMonitorV2,
+)
+from msprobe.core.monitor_v2.param import (
+    MindSporeParamMonitorV2,
+    PyTorchParamMonitorV2,
+)
+from msprobe.core.monitor_v2.weight_grad import (
+    MindSporeWeightGradMonitorV2,
+    PyTorchWeightGradMonitorV2,
+)
+
+MonitorCls = Type
+
+
+class MonitorFactory:
+    """
+    Registry-based factory that organizes monitors by framework, then by monitor name.
+    """
+
+    _REGISTRY: Dict[str, Dict[str, MonitorCls]] = {
+        "pytorch": {
+            "module": PyTorchModuleMonitorV2,
+            "weight_grad": PyTorchWeightGradMonitorV2,
+            "optimizer": PyTorchOptimizerMonitorV2,
+            "param": PyTorchParamMonitorV2,
+            "cc": PyTorchCCMonitorV2,
+        },
+        "mindspore": {
+            "module": MindSporeModuleMonitorV2,
+            "weight_grad": MindSporeWeightGradMonitorV2,
+            "optimizer": MindSporeOptimizerMonitorV2,
+            "param": MindSporeParamMonitorV2,
+            "cc": MindSporeCCMonitorV2,
+        },
+    }
+
+    @classmethod
+    def create(cls, framework: str, name: str) -> Optional[Any]:
+        fwk_key = str(framework).lower()
+        name_key = str(name).lower()
+        fwk = cls._REGISTRY.get(fwk_key)
+        if not fwk:
+            logger.warning(f"[monitor_v2] Unsupported framework '{framework}'.")
+            return None
+        mon_cls = fwk.get(name_key)
+        if mon_cls is None:
+            logger.warning(f"[monitor_v2] Unsupported monitor '{name}' for framework '{framework}'.")
+        return mon_cls() if mon_cls else None
+
+    @classmethod
+    def available(cls, framework: str) -> Dict[str, MonitorCls]:
+        return cls._REGISTRY.get(str(framework).lower(), {}).copy()

msprobe/core/monitor_v2/module.py

@@ -0,0 +1,201 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+
+from msprobe.core.common.const import Const, MonitorConst
+from msprobe.core.common.framework_adapter import FmkAdp
+from msprobe.core.common.log import logger
+from msprobe.core.monitor_v2.base import BaseMonitorV2
+from msprobe.core.monitor_v2.utils import get_vpp_stage_from_tag, iter_model_modules
+
+
+def _iter_named_modules(model: Any):
+    return iter_model_modules(model)
+
+
+class ModuleMonitorV2(BaseMonitorV2, ABC):
+    def __init__(self) -> None:
+        super().__init__()
+        self._model: Optional[Any] = None
+        self._hooks: List[Any] = []
+        self._targets: List[str] = []
+        self._eps: float = 1e-8
+
+    @staticmethod
+    def _try_register_hook(register_fn: Any, hook: Any) -> Any:
+        if register_fn is None:
+            return None
+        try:
+            return register_fn(hook, with_kwargs=True)
+        except TypeError:
+            return register_fn(hook)
+
+    def set_config(self, config: Dict[str, Any]) -> None:
+        super().set_config(config)
+        self.config = config
+        self._targets = config.get("targets", [])
+        self._eps = float(config.get("eps", 1e-8))
+
+    @abstractmethod
+    def _register_backward_hook(self, module: Any, module_name: str) -> Any:
+        ...
+
+    @abstractmethod
+    def _compute_metrics(self, tag2tensor: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
+        ...
+
+    def start(self, model: Any = None, optimizer: Any = None, **context: Any) -> None:
+        self.set_context(**context)
+        if self._model is not None:
+            return
+        model = model or self._context.get("model")
+        if model is None:
+            raise ValueError("Model must be provided to start ModuleMonitorV2")
+        self._model = model
+        self._register_hooks(model)
+
+    def stop(self) -> None:
+        if self._model is None:
+            return
+        for hook in self._hooks:
+            try:
+                hook.remove()
+            except Exception as exc:
+                logger.warning(f"[monitor_v2] Failed to remove module hook: {exc}")
+        self._hooks.clear()
+        self._model = None
+
+    def _register_hooks(self, model: Any) -> None:
+        targets = self._targets
+        for name, module in _iter_named_modules(model):
+            if not targets or any(keyword in name for keyword in targets):
+                self._register_module_hooks(module, name)
+
+    def _register_module_hooks(self, module: Any, module_name: str) -> None:
+        if not FmkAdp.is_nn_module(module):
+            return
+        self._add_hook(self._try_register_hook(
+            module.register_forward_pre_hook,
+            self._build_forward_pre_hook(module_name),
+        ))
+        self._add_hook(self._try_register_hook(
+            module.register_forward_hook,
+            self._build_forward_hook(module_name),
+        ))
+        backward_handle = self._register_backward_hook(module, module_name)
+        self._add_hook(backward_handle)
+
+    def _build_forward_pre_hook(self, module_name: str):
+        def hook(module: Any, inputs: Any, kwargs: Optional[Dict[str, Any]] = None):
+            tag2tensor, idx = self._collect_tagged_tensors(module_name, "input", inputs)
+            if kwargs:
+                extras, _ = self._collect_tagged_tensors(module_name, "input", tuple(kwargs.values()), idx)
+                tag2tensor.update(extras)
+            self._record_hook_data(module_name, "input", tag2tensor, "forward")
+        return hook
+
+    def _build_forward_hook(self, module_name: str):
+        def hook(module: Any, *hook_args: Any, **hook_kwargs: Any):
+            if len(hook_args) < 2:
+                return
+            outputs = hook_args[-1]
+            tag2tensor, _ = self._collect_tagged_tensors(module_name, "output", outputs)
+            self._record_hook_data(module_name, "output", tag2tensor, "forward")
+        return hook
+
+    def _build_backward_hook(self, module_name: str):
+        def hook(module: Any, grad_inputs: Any, grad_outputs: Any):
+            input_tags, _ = self._collect_tagged_tensors(module_name, "grad_input", grad_inputs)
+            output_tags, _ = self._collect_tagged_tensors(module_name, "grad_output", grad_outputs)
+            self._record_hook_data(module_name, "grad_input", input_tags, "backward")
+            self._record_hook_data(module_name, "grad_output", output_tags, "backward")
+        return hook
+
+    def _record_hook_data(self, module_name: str, io_kind: str, tag2tensor: Dict[str, Any], hook_type: str) -> None:
+        if not tag2tensor:
+            return
+        stats = self._compute_metrics(tag2tensor)
+        if not stats:
+            return
+        for tag, op_dict in stats.items():
+            vpp_stage = get_vpp_stage_from_tag(tag)
+            row = {
+                "module_name": tag,
+                "scope": hook_type,
+                "stats": {op: self._detach_stat(value) for op, value in op_dict.items()},
+            }
+            if vpp_stage is not None:
+                row["vpp_stage"] = vpp_stage
+            else:
+                row["vpp_stage"] = MonitorConst.DEFAULT_STAGE
+            self._rows.append(row)
+
+    def _collect_tagged_tensors(
+        self,
+        module_name: str,
+        io_kind: str,
+        tensors: Any,
+        start_idx: int = 0,
+    ) -> Tuple[Dict[str, Any], int]:
+        tagged: Dict[str, Any] = {}
+        idx = start_idx
+        if tensors is None:
+            return tagged, idx
+        if FmkAdp.is_tensor(tensors):
+            tagged[f"{module_name}.{io_kind}.{idx}"] = tensors
+            return tagged, idx + 1
+        if isinstance(tensors, (list, tuple)):
+            for tensor in tensors:
+                if not FmkAdp.is_tensor(tensor):
+                    continue
+                tagged[f"{module_name}.{io_kind}.{idx}"] = tensor
+                idx += 1
+        return tagged, idx
+
+    def _detach_stat(self, value: Any) -> Any:
+        if value is None or not FmkAdp.is_tensor(value):
+            return value
+        if FmkAdp.fmk == Const.PT_FRAMEWORK:
+            return value.detach()
+        return value
+
+    def _add_hook(self, handle: Any) -> None:
+        if handle is not None:
+            self._hooks.append(handle)
+
+
+class PyTorchModuleMonitorV2(ModuleMonitorV2):
+    def _register_backward_hook(self, module: Any, module_name: str) -> Any:
+        return module.register_full_backward_hook(self._build_backward_hook(module_name))
+
+    def _compute_metrics(self, tag2tensor: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
+        from msprobe.pytorch.monitor.module_metric import get_metrics
+
+        return get_metrics(self._ops, tag2tensor, self._eps, {})
+
+
+class MindSporeModuleMonitorV2(ModuleMonitorV2):
+    def _register_backward_hook(self, module: Any, module_name: str) -> Any:
+        return module.register_backward_hook(self._build_backward_hook(module_name))
+
+    def _compute_metrics(self, tag2tensor: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
+        from msprobe.mindspore.monitor.utils import get_metrics
+
+        return get_metrics(self._ops, tag2tensor, self._eps, {})