mindstudio-probe 1.0.1__py3-none-any.whl → 1.0.4__py3-none-any.whl

msprobe/mindspore/grad_probe/global_context.py

@@ -0,0 +1,90 @@
+import os
+import threading
+from typing import Dict, Union
+
+from msprobe.core.grad_probe.utils import check_str
+from msprobe.core.grad_probe.constant import GradConst
+from msprobe.mindspore.common.log import logger
+from msprobe.core.common.file_utils import create_directory, check_path_before_create
+
+
+class GlobalContext:
+
+    _instance = None
+    _instance_lock = threading.Lock()
+    _setting = {
+        GradConst.LEVEL: None,
+        GradConst.PARAM_LIST: None,
+        GradConst.STEP: None,
+        GradConst.RANK: None,
+        GradConst.CURRENT_STEP: 0,
+        GradConst.BOUNDS: [-10, -1, -0.1, -0.01, -0.001, 0, 0.001, 0.01, 0.1, 1, 10],
+        GradConst.OUTPUT_PATH: None
+    }
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance_lock.acquire()
+            cls._instance = object.__new__(cls)
+            cls._instance_lock.release()
+        return cls._instance
+
+    def init_context(self, config_dict: Dict):
+        level = config_dict.get(GradConst.LEVEL)
+        check_str(level, variable_name = "level in yaml")
+        if level in GradConst.SUPPORTED_LEVEL:
+            self._setting[GradConst.LEVEL] = config_dict.get(GradConst.LEVEL)
+        else:
+            raise ValueError("Invalid level set in config yaml file, level option: L0, L1, L2")
+
+        self._set_input_list(config_dict, GradConst.PARAM_LIST, str)
+        self._set_input_list(config_dict, GradConst.BOUNDS, float)
+        self._set_input_list(config_dict, GradConst.STEP, int)
+        self._set_input_list(config_dict, GradConst.RANK, int)
+
+        output_path = config_dict.get(GradConst.OUTPUT_PATH)
+        check_str(output_path, variable_name = "output_path in yaml")
+        try:
+            check_path_before_create(output_path)
+        except RuntimeError as err:
+            raise ValueError(f"Invalid output_path: {output_path}. The error message is {err}.") from err
+        self._setting[GradConst.OUTPUT_PATH] = output_path
+        if not os.path.isdir(self._setting.get(GradConst.OUTPUT_PATH)):
+            create_directory(self._setting.get(GradConst.OUTPUT_PATH))
+        else:
+            logger.warning("The output_path exists, the data will be covered.")
+
+    def get_context(self, key: str):
+        if key not in self._setting:
+            logger.warning(f"Unrecognized {key}.")
+        return self._setting.get(key)
+
+    def update_step(self):
+        self._setting[GradConst.CURRENT_STEP] += 1
+
+    def step_need_dump(self, step):
+        dump_step_list = self.get_context(GradConst.STEP)
+        return (not dump_step_list) or (step in dump_step_list)
+
+    def rank_need_dump(self, rank):
+        dump_rank_list = self.get_context(GradConst.RANK)
+        return (not dump_rank_list) or (rank in dump_rank_list)
+
+    def _set_input_list(self, config_dict: Dict, name: str, dtype: Union[int, str, float]):
+        value = config_dict.get(name)
+        if dtype == int:
+            type_str = "integer"
+        elif dtype == float:
+            type_str = "float"
+        else:
+            type_str = "string"
+        if value and isinstance(value, list):
+            for val in value:
+                if not isinstance(val, dtype):
+                    logger.warning(f"Invalid {name} which must be None or list of {type_str}")
+                    return
+            self._setting[name] = value
+        else:
+            logger.warning(f"{name} is None or not a list with valid items, use default value.")
+
+grad_context = GlobalContext()

msprobe/mindspore/grad_probe/grad_analyzer.py

@@ -0,0 +1,231 @@
+import os
+import time
+from typing import List, Tuple
+import multiprocessing
+from multiprocessing import Process
+
+import numpy as np
+import mindspore as ms
+from mindspore.communication import get_rank
+from mindspore.ops import operations as P
+from mindspore.common.parameter import Parameter
+
+from msprobe.core.grad_probe.utils import ListCache
+from msprobe.core.grad_probe.constant import GradConst
+from msprobe.mindspore.common.log import logger
+from msprobe.core.common.file_utils import (create_directory, check_file_or_directory_path,
+                                            write_csv, remove_path, move_file, load_npy)
+from msprobe.mindspore.grad_probe.global_context import grad_context, GlobalContext
+
+
+def get_rank_id():
+    try:
+        rank_id = get_rank()
+    except Exception as err:
+        rank_id = 0
+    return rank_id
+
+
+@ms.jit
+def grad_dump(dump_dir: str, g_name: str, dump_step: Parameter, grad: ms.Tensor, level: str, bounds: List):
+    '''
+    Dump gradient statistic data.
+        level0: [step, max, min, norm, shape_dim, shape]
+        level1: [step, max, min, norm, shape_dim, shape] + grad_bool_data
+        level2: [step, max, min, norm, shape_dim, shape, dist_dim, dist] + grad_bool_data
+    '''
+    dump_path = os.path.join(dump_dir, g_name)
+    dump_dir_path = dump_path + "_dir"
+    save_op = ms.ops.TensorDump()
+
+    grad_flat = grad.reshape(-1)
+    max_val = grad_flat.max(axis=0).float()
+    min_val = grad_flat.min(axis=0).float()
+    norm_val = grad_flat.norm(ord=2).float()
+    shape = grad.shape
+    extrem_list = [dump_step[0].float(), max_val, min_val, norm_val]
+    extrem_stat = ms.ops.stack(extrem_list)
+    shape_list = [len(shape)] + list(shape)
+    shape_stat = ms.Tensor(shape_list).float()
+    level0_stat = ms.ops.concat((extrem_stat, shape_stat), axis=0)
+    level_stat = level0_stat
+
+    if level == GradConst.LEVEL2:
+        zero_grad = (grad == 0).sum()
+        dist_dim = ms.Tensor([len(bounds) + 2]).float()
+        bucket_result = ms.ops.bucketize(grad.float(), bounds)
+        bucket_result = bucket_result.astype(ms.int8)
+        dist_stat = [(bucket_result == i).sum() for i in range(len(bounds) + 1)]
+        dist_stat.append(zero_grad)
+        dist_stat.append(ms.Tensor(1, dtype=ms.int64))  # make sure dist_stat is not empty
+        dist_stat = ms.ops.stack(dist_stat, axis=0).float()
+        level2_stat = ms.ops.concat((level0_stat, dist_dim, dist_stat), axis=0)
+        level_stat = level2_stat
+
+    save_op(dump_path, level_stat)
+    if level == GradConst.LEVEL1 or level == GradConst.LEVEL2:
+        grad_direction = grad > 0
+        save_op(dump_dir_path, grad_direction)
+
+
+class CSVGenerator(Process):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dump_dir = None
+        self.save_dir = None
+        self.level = GradConst.LEVEL0
+        self.cache_list = ListCache()
+        self.current_step = None
+        self.stop_event = None
+        self.last_finish = False
+        self.bounds = [-0.1, 0.0, 0.1],
+
+    def init(self, context: GlobalContext):
+        rank_id = get_rank_id()
+        output_path = context.get_context(GradConst.OUTPUT_PATH)
+        self.level = context.get_context(GradConst.LEVEL)
+        self.bounds = context.get_context(GradConst.BOUNDS)
+        self.dump_dir = f"{output_path}/rank{rank_id}/Dump/"
+        self.save_dir = f"{output_path}/rank{rank_id}/"
+        self.current_step = None
+        self.stop_event = multiprocessing.Event()
+        self.last_finish = False
+
+    def run(self):
+        while True:
+            if not os.path.exists(self.dump_dir):
+                time.sleep(0.1)
+                if self.stop_event.is_set():
+                    break
+                continue
+            npy_files = os.listdir(self.dump_dir)
+            npy_files.sort(key=lambda x: int(x.split("_")[0]))
+            self.traverse_files(npy_files)
+            empty = len(os.listdir(self.dump_dir)) == 0
+            if self.stop_event.is_set() and empty and self.last_finish:
+                break
+        if os.path.exists(self.dump_dir):
+            remove_path(self.dump_dir)
+
+    def stop(self):
+        self.stop_event.set()
+
+    def traverse_files(self, npy_files: List):
+        for npy_file in npy_files:
+            file_path = os.path.join(self.dump_dir, npy_file)
+            while not os.path.exists(file_path):
+                time.sleep(0.01)
+            check_file_or_directory_path(file_path)
+            if GradConst.STEP_FINISH in npy_file:
+                self.cache_list.flush()
+                remove_path(file_path)
+                self.last_finish = True
+            elif file_path.split("_")[-1] == GradConst.DIR_SUFFIX:
+                prefix_idx = len(npy_file.split("_")[0])
+                new_name = npy_file[prefix_idx + 1:].replace("_" + GradConst.DIR_SUFFIX, "." + GradConst.NPY_SUFFIX)
+                if not new_name:
+                    raise RuntimeError("Invalid dump data name.")
+                if self.current_step is None:
+                    raise RuntimeError("Current record step is None.")
+                step_dir = os.path.join(self.save_dir, f"step{self.current_step}")
+                if not os.path.exists(step_dir):
+                    create_directory(step_dir)
+                dst_file = os.path.join(step_dir, new_name)
+                move_file(file_path, dst_file)
+                self.last_finish = False
+            elif file_path.split(".")[-1] == GradConst.NPY_SUFFIX:
+                stat_data = self.load_npy_data(file_path)
+                if stat_data is None:
+                    continue
+                if not self.check_valid(stat_data):
+                    remove_path(file_path)
+                    continue
+                step = int(stat_data[GradConst.STEP_IDX])
+                update_step = self.current_step is None or step != self.current_step
+                self.current_step = step
+                if update_step:
+                    self.create_csv_file()
+                self.gen_csv_line(file_path, stat_data)
+                remove_path(file_path)
+                self.last_finish = False
+
+    def check_valid(self, stat_data):
+        level = grad_context.get_context(GradConst.LEVEL)
+        try:
+            shape_dim = int(stat_data[GradConst.SHAPE_DIM_IDX])
+            if level == GradConst.LEVEL2:
+                dist_dim = int(stat_data[shape_dim + GradConst.SHAPE_DIM_IDX + 1])
+                length = shape_dim + dist_dim + 7
+            else:
+                length = shape_dim + 5
+        except IndexError as err:
+            return False
+        if length != len(stat_data):
+            return False
+        return True
+
+    def load_npy_data(self, file_path: str):
+        stat_data = None
+        max_try = 10
+        while max_try:
+            try:
+                stat_data = load_npy(file_path)
+                return stat_data
+            except Exception as err:
+                logger.warning(f"load numpy file failed, retry...")
+                max_try -= 1
+                time.sleep(0.1)
+        return stat_data
+
+    def gen_csv_line(self, file_path: str, stat_data) -> None:
+        shape_dim = int(stat_data[GradConst.SHAPE_DIM_IDX])
+        file_name = os.path.basename(file_path)
+        prefix_idx = len(file_name.split("_")[0])
+        param_name = file_name[(prefix_idx + 1) : -(len(GradConst.NPY_SUFFIX) + 1)]
+        if not param_name:
+            raise RuntimeError("Invalid gradient statistic file name.")
+        csv_line = [param_name]
+        if self.level == GradConst.LEVEL2:
+            csv_line.extend(self.get_dist_data(shape_dim, stat_data))
+        csv_line.extend(self.get_extrem_data(shape_dim, stat_data))
+        self.cache_list.append(csv_line)
+
+    def get_dist_data(self, shape_dim: int, stat_data: np.ndarray):
+        dist_data = stat_data[(shape_dim + GradConst.SHAPE_DIM_IDX + 2):-1]
+        element_num = dist_data.sum() - dist_data[-1]
+        if element_num != 0:
+            dist_data = dist_data / element_num
+        return list(dist_data)
+
+    def get_extrem_data(self, shape_dim: int, stat_data: np.ndarray):
+        extrem_data = list(stat_data[(GradConst.STEP_IDX + 1):(GradConst.STEP_IDX + 4)])
+        shape_data = stat_data[(GradConst.SHAPE_DIM_IDX + 1):(GradConst.SHAPE_DIM_IDX + shape_dim + 1)]
+        shape_data = list(shape_data.astype(int))
+        extrem_data.append(shape_data)
+        return extrem_data
+
+    def create_csv_file(self):
+        headers = ["Param_name"]
+        if self.level == GradConst.LEVEL2:
+            headers.extend(self.get_dist_header())
+        headers.extend(self.get_extrem_headers())
+        output_path = f"{self.save_dir}/grad_summary_{self.current_step}.csv"
+        write_csv([headers], output_path)
+        self.cache_list.set_output_file(output_path)
+        self.cache_list.clear()
+
+    def get_extrem_headers(self) -> List[str]:
+        return ["Max", "Min", "Norm", "Shape"]
+
+    def get_dist_header(self) -> List[str]:
+        intervals = []
+        for i, _ in enumerate(self.bounds):
+            if i == 0:
+                intervals.append(f"(-inf, {self.bounds[i]}]")
+            else:
+                intervals.append(f"({self.bounds[i-1]}, {self.bounds[i]}]")
+        intervals.extend([f"({self.bounds[-1]}, inf)", "=0"])
+        return intervals
+
+csv_generator = CSVGenerator()

msprobe/mindspore/grad_probe/grad_monitor.py

@@ -0,0 +1,27 @@
+from msprobe.mindspore.grad_probe.global_context import grad_context
+from msprobe.mindspore.grad_probe.grad_analyzer import csv_generator
+from msprobe.mindspore.grad_probe.hook import hook_optimizer
+from msprobe.core.grad_probe.constant import GradConst
+
+
+class GradientMonitor:
+
+    def __init__(self, common_dict, task_config):
+        config = {}
+        config[GradConst.OUTPUT_PATH] = common_dict.dump_path
+        config[GradConst.STEP] = common_dict.step
+        config[GradConst.RANK] = common_dict.rank
+        config[GradConst.PARAM_LIST] = task_config.param_list
+        config[GradConst.LEVEL] = task_config.grad_level
+        config[GradConst.BOUNDS] = task_config.bounds
+        self.config = config
+        grad_context.init_context(self.config)
+
+    @staticmethod
+    def monitor(opt):
+        csv_generator.init(grad_context)
+        hook_optimizer(opt)
+
+    @staticmethod
+    def stop():
+        csv_generator.stop()

msprobe/mindspore/grad_probe/grad_stat_csv.py

@@ -0,0 +1,132 @@
+from abc import ABC, abstractmethod
+import hashlib
+
+import mindspore
+from mindspore import ops, Tensor
+from msprobe.core.grad_probe.constant import GradConst
+
+
+class CsvInput:
+    def __init__(self, param_name, grad, bounds):
+        self.param_name = param_name
+        self.grad = grad
+        self.bounds = bounds
+
+class GradStatCsv:
+    csv = {}
+
+    @staticmethod
+    def get_csv_header(level, csv_input):
+        header = ["param_name"]
+        for key in level["header"]:
+            header.extend(GradStatCsv.csv[key].generate_csv_header(csv_input))
+        return header
+
+    @staticmethod
+    def get_csv_line(level, csv_input):
+        line = [csv_input.param_name]
+        for key in level["header"]:
+            line.extend(GradStatCsv.csv[key].generate_csv_content(csv_input))
+        return line
+
+
+def register_csv_item(key, cls=None):
+    if cls is None:
+        # 无参数时，返回装饰器函数
+        return lambda cls: register_csv_item(key, cls)
+    GradStatCsv.csv[key] = cls
+    return cls
+
+
+class CsvItem(ABC):
+    @staticmethod
+    @abstractmethod
+    def generate_csv_header(csv_input):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def generate_csv_content(csv_input):
+        pass
+
+
+@register_csv_item(GradConst.MD5)
+class CsvMd5(CsvItem):
+    def generate_csv_header(csv_input):
+        return ["MD5"]
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        tensor_bytes = grad.float().numpy().tobytes()
+        md5_hash = hashlib.md5(tensor_bytes)
+        return [md5_hash.hexdigest()]
+
+
+@register_csv_item(GradConst.DISTRIBUTION)
+class CsvDistribution(CsvItem):
+    def generate_csv_header(csv_input):
+        bounds = csv_input.bounds
+        intervals = []
+        if bounds:
+            intervals.append(f"(-inf, {bounds[0]}]")
+            for i in range(1, len(bounds)):
+                intervals.append(f"({bounds[i-1]}, {bounds[i]}]")
+        if intervals:
+            intervals.append(f"({bounds[-1]}, inf)")
+        intervals.append("=0")
+    
+        return intervals
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        bounds = csv_input.bounds
+        if grad.dtype == mindspore.bfloat16:
+            grad = grad.to(mindspore.float32)
+        element_num = grad.numel()
+        grad_equal_0_num = (grad == 0).sum().item()
+        bucketsize_result = ops.bucketize(grad.float(), bounds)
+        bucketsize_result = bucketsize_result.astype(mindspore.int8)
+        interval_nums = [(bucketsize_result == i).sum().item() for i in range(len(bounds) + 1)]
+        interval_nums.append(grad_equal_0_num)
+        return_list = [x / element_num if element_num != 0 else 0 for x in interval_nums]
+        return return_list
+
+
+@register_csv_item(GradConst.MAX)
+class CsvMax(CsvItem):
+    def generate_csv_header(csv_input):
+        return ["max"]
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        return [ops.amax(grad).float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.MIN)
+class CsvMin(CsvItem):
+    def generate_csv_header(csv_input):
+        return ["min"]
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        return [ops.amin(grad).float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.NORM)
+class CsvNorm(CsvItem):
+    def generate_csv_header(csv_input):
+        return ["norm"]
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        return [ops.norm(grad).float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.SHAPE)
+class CsvShape(CsvItem):
+    def generate_csv_header(csv_input):
+        return ["shape"]
+
+    def generate_csv_content(csv_input):
+        grad = csv_input.grad
+        return [list(grad.shape)]

msprobe/mindspore/grad_probe/hook.py

@@ -0,0 +1,94 @@
+
+import os
+
+import mindspore
+import mindspore as ms
+from mindspore.common.api import jit
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.common.parameter import Parameter
+from mindspore.common.initializer import initializer
+
+from msprobe.core.grad_probe.constant import GradConst
+from msprobe.mindspore.common.log import logger
+
+from msprobe.core.common.file_utils import remove_path, write_csv, create_directory
+from msprobe.mindspore.grad_probe.global_context import grad_context
+from msprobe.mindspore.grad_probe.grad_analyzer import grad_dump, get_rank_id
+from msprobe.mindspore.grad_probe.grad_analyzer import csv_generator
+from msprobe.mindspore.grad_probe.grad_stat_csv import GradStatCsv, CsvInput
+from msprobe.mindspore.grad_probe.utils import save_grad_direction, get_adapted_level
+
+class HookInput:
+
+    '''
+    HookInput is a class wrapping all the variables used for hooking optimizer
+    '''
+
+    def __init__(self, opt) -> None:
+        self.func = opt.construct
+        self.g_names = [param.name for param in opt._parameters]
+        self.param_list = grad_context.get_context(GradConst.PARAM_LIST)
+        self.rank_id = get_rank_id()
+        output_path = grad_context.get_context(GradConst.OUTPUT_PATH)
+        self.dump_dir = os.path.join(output_path, f"rank{self.rank_id}", "Dump")
+        self.save_dir = os.path.join(output_path, f"rank{self.rank_id}")
+        self.step_finish_flag = os.path.join(self.dump_dir, GradConst.STEP_FINISH)
+        if os.path.exists(self.save_dir):
+            logger.warning(f"Delete existing path {self.save_dir}.")
+            remove_path(self.save_dir)
+        self.level = grad_context.get_context(GradConst.LEVEL)
+        self.bounds = grad_context.get_context(GradConst.BOUNDS)
+        self.mode = mindspore.get_context("mode")
+
+def hook_graph_mode_optimizer(opt, hook_input):
+    @jit
+    def new_construct(self, gradients):
+        for index, grad_value in enumerate(gradients):
+            if hook_input.param_list and hook_input.g_names[index] not in hook_input.param_list:
+                continue
+            grad_dump(hook_input.dump_dir, hook_input.g_names[index], self.dump_step,
+                    grad_value, hook_input.level, hook_input.bounds)
+        ms.ops.TensorDump()(hook_input.step_finish_flag, self.dump_step)
+        self.assignadd(self.dump_step, self.global_step_increase_tensor)
+        out = hook_input.func(gradients)
+        return out
+
+    opt.dump_step = Parameter(initializer(0, [1], ms.int32), name="dump_step")
+    opt.construct = new_construct.__get__(opt, type(opt))
+    csv_generator.start()
+
+def hook_pynative_optimizer(opt, hook_input):
+    level_adapted = get_adapted_level(hook_input.level)
+
+    def hook_fn(cell, input):
+        gradients, = input
+        cur_step = grad_context.get_context(GradConst.CURRENT_STEP)
+        if grad_context.step_need_dump(cur_step) and grad_context.rank_need_dump(hook_input.rank_id):
+            create_directory(hook_input.save_dir)
+            output_lines = []
+            for index, grad_value in enumerate(gradients):
+                param_name = hook_input.g_names[index]
+                if hook_input.param_list and param_name not in hook_input.param_list:
+                    continue
+                csv_input = CsvInput(param_name, grad_value, hook_input.bounds)
+                grad_info = GradStatCsv.get_csv_line(level_adapted, csv_input)
+                output_lines.append(grad_info)
+                if level_adapted["have_grad_direction"]:
+                    save_grad_direction(param_name, grad_value, os.path.join(hook_input.save_dir, f'step{cur_step}'))
+            output_csv_path = os.path.join(hook_input.save_dir, f"grad_summary_{cur_step}.csv")
+            dummy_csv_input = CsvInput(None, None, hook_input.bounds)
+            output_lines.insert(0, GradStatCsv.get_csv_header(level_adapted, dummy_csv_input))
+            write_csv(output_lines, output_csv_path)
+            logger.info(f"write grad data to {output_csv_path}")
+        grad_context.update_step()
+
+    opt.register_forward_pre_hook(hook_fn)
+
+
+def hook_optimizer(opt: Optimizer):
+    hook_input = HookInput(opt)
+
+    if hook_input.mode == mindspore.GRAPH_MODE:
+        hook_graph_mode_optimizer(opt, hook_input)
+    else:
+        hook_pynative_optimizer(opt, hook_input)

msprobe/mindspore/grad_probe/utils.py

@@ -0,0 +1,30 @@
+import os
+
+import mindspore
+from msprobe.core.grad_probe.constant import level_adp
+from msprobe.core.grad_probe.utils import check_param
+from msprobe.core.common.file_utils import (create_directory,
+                                            check_path_before_create,
+                                            check_file_or_directory_path,
+                                            save_npy)
+
+
+def save_grad_direction(param_name, grad, save_path):
+    if not os.path.exists(save_path):
+        create_directory(save_path)
+    check_file_or_directory_path(save_path, isdir=True)
+    check_param(param_name)
+    save_filepath = os.path.join(save_path, f"{param_name}.npy")
+    check_path_before_create(save_filepath)
+
+    if grad.dtype == mindspore.bfloat16:
+        grad = grad.to(mindspore.float32)
+    grad_direction_tensor = grad > 0
+    grad_direction_ndarray = grad_direction_tensor.numpy()
+
+    save_npy(grad_direction_ndarray, save_filepath)
+
+
+def get_adapted_level(level: str):
+    level_adapted = level_adp.get(level)
+    return level_adapted