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

msprobe/core/monitor_v2/optimizer.py

@@ -0,0 +1,245 @@
+# -------------------------------------------------------------------------
+#  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
+
+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 build_param2name, get_vpp_stage_from_tag
+
+
+class OptimizerMonitorV2(BaseMonitorV2):
+    """
+    通用 Optimizer 监控基类（当前只实现 mv_distribution，后续可扩展 wg/param 等）。
+
+    公共职责：
+      - 解析 config（eps / mv_distribution / ops）
+      - 记录 model / optimizer / context
+      - 基于模型构建 param -> name 映射
+      - 将 (exp_avg_dict, exp_avg_sq_dict) 转成 rows（module_name / hook_type / stats）
+
+    框架差异由子类通过以下方法提供：
+      - _build_backend
+      - _get_mv
+      - _compute_metrics
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._optimizer: Any = None
+        self._backend: Any = None
+        self._param2name: Dict[Any, str] = {}
+        self._eps: float = 1e-8
+
+        self.mv_distribution: bool = True
+        self.mg_direction: bool = False
+        self.ur_distribution: bool = False
+
+    def set_config(self, config: Dict[str, Any]) -> None:
+        super().set_config(config)
+        self.config = config
+        self._eps = float(config.get("eps", 1e-8))
+        self.mv_distribution = bool(config.get("mv_distribution", True))
+
+    def start(self, model: Any = None, optimizer: Any = None, **context: Any) -> None:
+        self.set_context(**context)
+        if self._optimizer is not None:
+            return
+        optimizer = optimizer or self._context.get("optimizer")
+        model = model or self._context.get("model")
+        if optimizer is None or model is None:
+            raise ValueError("Model and optimizer must be provided to start optimizer function")
+        self._optimizer = optimizer
+        self._register_params_v2(model)
+        self._backend = self._build_backend(optimizer)
+
+    def stop(self) -> None:
+        self._optimizer = None
+        self._backend = None
+        self._param2name.clear()
+        self._rows = []
+
+    def collect(self) -> Dict[str, Any] | None:
+        if not self.mv_distribution or not self._param2name:
+            return None
+        exp_avg_dict, exp_avg_sq_dict = self._get_mv()
+        if not exp_avg_dict and not exp_avg_sq_dict:
+            return None
+        self._rows = []
+        self._record_mv("exp_avg", exp_avg_dict)
+        self._record_mv("exp_avg_sq", exp_avg_sq_dict)
+        return super().collect()
+
+    def _register_params_v2(self, model: Any) -> None:
+        """
+        使用统一工具函数建立 param->name 映射，支持
+        单个 model 或 model 列表。
+        """
+        self._param2name = build_param2name(model=model, optimizer=self._optimizer)
+
+    def _register_params(self, model: Any) -> None:
+        """
+        默认使用模型的 named_parameters 建立 param->name 映射。
+        子类可在特殊场景中覆写。
+        """
+        self._param2name.clear()
+        try:
+            for name, param in FmkAdp.named_parameters(model):
+                self._param2name[param] = name
+        except Exception:
+            self._param2name.clear()
+
+    def _record_mv(self, metric_type: str, name2tensor: Dict[str, Any]) -> None:
+        if not name2tensor:
+            return
+        tag2tensor: Dict[str, Any] = {}
+        for name, tensor in name2tensor.items():
+            if tensor is None or not FmkAdp.is_tensor(tensor):
+                continue
+            tag2tensor[f"{name}.{metric_type}"] = tensor
+        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": metric_type,
+                "stats": {op: self._detach_stat(val) for op, val 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 _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 _build_backend(self, optimizer: Any) -> Any:  # pragma: no cover - abstract by convention
+        raise NotImplementedError
+
+    def _get_mv(self) -> tuple[Dict[str, Any], Dict[str, Any]]:  # pragma: no cover - abstract by convention
+        raise NotImplementedError
+
+    def _compute_metrics(self, tag2tensor: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:  # pragma: no cover
+        raise NotImplementedError
+
+
+class PyTorchOptimizerMonitorV2(OptimizerMonitorV2):
+    """
+    PyTorch optimizer monitor，当前实现 mv_distribution 部分。
+    """
+
+    def _build_backend(self, optimizer: Any) -> Any:
+        from msprobe.pytorch.monitor.optimizer_collect import OptimizerMonFactory
+
+        return OptimizerMonFactory.create_optimizer_mon(optimizer)
+
+    def _get_mv(self) -> tuple[Dict[str, Any], Dict[str, Any]]:
+        from msprobe.core.monitor.utils import MVResult
+
+        if self._backend is None:
+            return {}, {}
+        result = self._backend.fetch_mv(self, self._param2name)
+        if isinstance(result, MVResult):
+            return result.exp_avg, result.exp_avg_sq
+        try:
+            exp_avg_dict, exp_avg_sq_dict = result[0], result[1]
+            return exp_avg_dict, exp_avg_sq_dict
+        except Exception:
+            return {}, {}
+
+    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 MindSporeOptimizerMonitorV2(OptimizerMonitorV2):
+    """
+    MindSpore optimizer monitor，当前实现 mv_distribution 部分。
+
+    对齐 mindspore/monitor 中的两条路径：
+      - mindtorch 场景：使用 optimizer_collect.OptimizerMon.fetch_mv
+      - 纯 MindSpore 场景：使用 get_mv_for_ms 从 optimizer.parameters 中抓取 m/v
+    """
+
+    def _build_backend(self, optimizer: Any) -> Any:
+        from msprobe.mindspore.monitor.optimizer_collect import OptimizerMonFactory
+
+        return OptimizerMonFactory.create_optimizer_mon(optimizer)
+
+    def _get_mv(self) -> tuple[Dict[str, Any], Dict[str, Any]]:
+        from msprobe.mindspore.common.utils import is_mindtorch
+
+        if self._optimizer is None:
+            return {}, {}
+        if is_mindtorch():
+            if self._backend is None:
+                return {}, {}
+            result = self._backend.fetch_mv(self, self._param2name)
+            try:
+                exp_avg_dict, exp_avg_sq_dict = result[0], result[1]
+            except Exception:
+                return {}, {}
+            return exp_avg_dict, exp_avg_sq_dict
+        # 纯 MindSpore 场景：使用原生 get_mv_for_ms 语义
+        return self._get_mv_for_ms(self._optimizer)
+
+    def _get_mv_for_ms(self, opt: Any) -> tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        直接复用 mindspore/monitor.module_hook.get_mv_for_ms 的采集语义：
+          - 若是纯 MindSpore，optimizer 可能包了一层，需要解出真正的 Cell
+          - 从 get_parameters(common_opt) 中抓取名字带 EXP_AVG / EXP_AVG_SQ 的参数
+        """
+        from msprobe.mindspore.monitor.common_func import is_valid_instance, get_parameters
+
+        if not self.mv_distribution:
+            return {}, {}
+
+        common_opt = opt
+        if not is_valid_instance(opt):
+            common_opt = getattr(opt, "optimizer", None)
+            if not is_valid_instance(common_opt):
+                logger.warning("MindSporeOptimizerMonitorV2: optimizer is not valid, please check usage")
+                return {}, {}
+
+        m_dict: Dict[str, Any] = {}
+        v_dict: Dict[str, Any] = {}
+        for name, param in get_parameters(common_opt):
+            if MonitorConst.EXP_AVG_SQ in name:
+                v_dict[name] = param
+            elif MonitorConst.EXP_AVG in name:
+                m_dict[name] = param
+
+        return m_dict, v_dict
+
+    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, {})

msprobe/core/monitor_v2/param.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+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 build_param2name, get_vpp_stage_from_tag
+
+
+class ParamMonitorV2(BaseMonitorV2, ABC):
+    """
+    Parameter distribution monitor that captures params before/after optimizer step.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._optimizer: Any = None
+        self._model: Any = None
+        self._param2name: Dict[Any, str] = {}
+        self._eps: float = 1e-8
+        self.param_distribution: bool = True
+        self._step_hook_handles: List[Any] = []
+        self._orig_step = None
+        self._step_patched = False
+
+    @abstractmethod
+    def _compute_metrics(self, tag2tensor: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
+        ...
+
+    def set_config(self, config: Dict[str, Any]) -> None:
+        super().set_config(config)
+        self.config = config
+        self._eps = float(config.get("eps", 1e-8))
+        self.param_distribution = bool(config.get("param_distribution", True))
+
+    def start(self, model: Any = None, optimizer: Any = None, **context: Any) -> None:
+        self.set_context(**context)
+        if self._optimizer is not None:
+            return
+        optimizer = optimizer or self._context.get("optimizer")
+        model = model or self._context.get("model")
+        if optimizer is None or model is None:
+            raise ValueError("Model and optimizer must be provided to start ParamMonitorV2")
+        self._optimizer = optimizer
+        self._model = model
+        self._param2name = build_param2name(model=model, optimizer=optimizer)
+        self._patch_optimizer()
+
+    def stop(self) -> None:
+        self._restore_optimizer()
+        self._optimizer = None
+        self._model = None
+        self._param2name.clear()
+        self._rows = []
+
+    def _patch_optimizer(self) -> None:
+        if not self.param_distribution or self._optimizer is None:
+            return
+        pre_hook = getattr(self._optimizer, "register_step_pre_hook", None)
+        post_hook = getattr(self._optimizer, "register_step_post_hook", None)
+        if callable(pre_hook) and callable(post_hook):
+            self._step_hook_handles.append(pre_hook(self._on_pre_step))
+            self._step_hook_handles.append(post_hook(self._on_post_step))
+            return
+        if self._step_patched or not hasattr(self._optimizer, "step"):
+            return
+        self._orig_step = self._optimizer.step
+
+        def patched_step(*args: Any, **kwargs: Any):
+            self._record_param_distribution(MonitorConst.PRE_PARAM)
+            out = self._orig_step(*args, **kwargs)
+            self._record_param_distribution(MonitorConst.POST_PARAM)
+            return out
+
+        self._optimizer.step = patched_step  # type: ignore[assignment]
+        self._step_patched = True
+
+    def _restore_optimizer(self) -> None:
+        for handle in self._step_hook_handles:
+            try:
+                handle.remove()
+            except Exception:
+                pass
+        self._step_hook_handles.clear()
+        if self._step_patched and self._optimizer is not None and self._orig_step is not None:
+            self._optimizer.step = self._orig_step  # type: ignore[assignment]
+        self._orig_step = None
+        self._step_patched = False
+
+    def _on_pre_step(self, optimizer: Any, args: Any, kwargs: Any) -> None:
+        self._record_param_distribution(MonitorConst.PRE_PARAM)
+
+    def _on_post_step(self, optimizer: Any, args: Any, kwargs: Any) -> None:
+        self._record_param_distribution(MonitorConst.POST_PARAM)
+
+    def _record_param_distribution(self, scope: str) -> None:
+        if not self.param_distribution or not self._param2name:
+            return
+        tag2tensor: Dict[str, Any] = {}
+        for param, name in self._param2name.items():
+            try:
+                if param.numel() == 0:
+                    continue
+            except Exception:
+                continue
+            if not FmkAdp.is_tensor(param):
+                continue
+            tensor = param
+            if tensor.__class__.__name__ == "DTensor":
+                try:
+                    tensor = tensor.to_local()
+                except Exception as exc:
+                    logger.warning(f"[monitor_v2] Failed to convert DTensor param to local: {exc}")
+                    continue
+            tag2tensor[name] = tensor
+        if not tag2tensor:
+            return
+        stats = self._compute_metrics(tag2tensor)
+        for tag, op_dict in stats.items():
+            vpp_stage = get_vpp_stage_from_tag(tag)
+            row = {
+                "module_name": tag,
+                "scope": scope,
+                "stats": {op: self._detach_stat(val) for op, val 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 _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
+
+
+class PyTorchParamMonitorV2(ParamMonitorV2):
+    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 MindSporeParamMonitorV2(ParamMonitorV2):
+    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, {})