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

msprobe/core/compare/utils.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");
@@ -21,7 +21,7 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from msprobe.core.common.const import Const, CompareConst
+from msprobe.core.common.const import Const, CompareConst, FileCheckConst
 from msprobe.core.common.utils import CompareException, check_regex_prefix_format_valid, logger, safe_get_value
 from msprobe.core.common.file_utils import check_file_or_directory_path
 
@@ -37,13 +37,20 @@ def extract_json(dirname, stack_json=False):
     # Provide robustness on invalid directory inputs
     if not json_path:
         if stack_json:
-            logger.error(f'stack.json is not found in dump dir {dirname}.')
+            logger.warning(f'stack.json is not found in dump dir {dirname}.')
         else:
             logger.error(f'dump.json is not found in dump dir {dirname}.')
-        raise CompareException(CompareException.NO_DUMP_FILE_ERROR)
+            raise CompareException(CompareException.NO_DUMP_FILE_ERROR)
     return json_path
 
 
+def set_stack_json_path(input_param):
+    npu_data_dir = os.path.dirname(input_param.get("npu_json_path"))
+    stack_path = extract_json(npu_data_dir, stack_json=True)
+    input_param["stack_json_path"] = stack_path if stack_path else None
+    return bool(stack_path)
+
+
 def check_and_return_dir_contents(dump_dir, prefix):
     """
     check the given dump dir and validate files in dump dir by using the given prefix patterns to build a
@@ -75,6 +82,10 @@ def check_and_return_dir_contents(dump_dir, prefix):
 
 
 def rename_api(npu_name, process):
+    """
+    原api： {api_type}.{api_name}.{API调用次数}.{前向反向}.{input/output}.{参数序号}
+    rename后： {api_type}.{api_name}.{input/output}.{参数序号}
+    """
     npu_split = npu_name.split(process)
     try:
         torch_func_index, in_out = npu_split[0], npu_split[1]
@@ -87,17 +98,13 @@ def rename_api(npu_name, process):
 
 
 def read_op(op_data, op_name):
-    io_name_mapping = {
-        Const.INPUT_ARGS: '.input',
-        Const.INPUT_KWARGS: '.input',
-        Const.INPUT: '.input',
-        Const.OUTPUT: '.output'
-    }
-
-    op_parsed_list = []
-    for name in io_name_mapping:
-        if name in op_data:
-            op_parsed_list.extend(op_item_parse(op_data[name], op_name + io_name_mapping[name]))
+    if Const.PARAMS_GRAD in op_name.split(Const.SEP):
+        op_parsed_list = op_item_parse(op_data, op_name)
+    else:
+        op_parsed_list = []
+        for name in CompareConst.IO_NAME_MAPPING:
+            if name in op_data:
+                op_parsed_list.extend(op_item_parse(op_data[name], op_name + CompareConst.IO_NAME_MAPPING[name]))
     return op_parsed_list
 
 
@@ -124,11 +131,14 @@ def op_item_parse(op_data, op_name: str, depth: int = 0) -> list:
         return [default_item]
     elif not op_data:
         return []
-    
+
     item_list = []
     if isinstance(op_data, list):
         for i, data in enumerate(op_data):
-            item_list.extend(op_item_parse(data, op_name + Const.SEP + str(i), depth + 1))
+            if Const.PARAMS_GRAD not in op_name.split(Const.SEP):
+                item_list.extend(op_item_parse(data, op_name + Const.SEP + str(i), depth + 1))
+            else:
+                item_list.extend(op_item_parse(data, op_name, depth + 1))
     elif isinstance(op_data, dict):
         if is_leaf_data(op_data):
             return [gen_op_item(op_data, op_name)]
@@ -144,14 +154,15 @@ def is_leaf_data(op_data):
 def gen_op_item(op_data, op_name):
     op_item = {}
     op_item.update(op_data)
-    op_item['full_op_name'] = op_name
-    op_item['data_name'] = op_data.get('data_name', '-1')
+    data_name = op_data.get('data_name') if op_data.get('data_name') else '-1'  # 如果是""也返回-1
+    op_item['data_name'] = data_name
+    op_item['full_op_name'] = data_name.rsplit(Const.SEP, 1)[0] if data_name != '-1' else op_name
 
     params = ['Max', 'Min', 'Mean', 'Norm']
     for i in params:
         if i not in op_item:
             op_item[i] = None
-    
+
     if not op_item.get('dtype'):
         if op_item.get('type') == 'torch.Size':
             op_item['dtype'] = op_data.get('type')
@@ -159,6 +170,16 @@ def gen_op_item(op_data, op_name):
         elif op_item.get('type') == 'slice':
             op_item['dtype'] = op_data.get('type')
             op_item['shape'] = str(np.shape(np.array(op_data.get('value'))))
+        elif op_item.get('type') == 'ellipsis':
+            op_item['dtype'] = op_data.get('type')
+            op_item['shape'] = '[]'
+            for i in params:
+                op_item[i] = op_data.get('value')
+        elif op_item.get('type') == 'torch.ProcessGroup':
+            op_item['dtype'] = op_data.get('type')
+            op_item['shape'] = '[]'
+            for i in params:
+                op_item[i] = str(op_data.get('group_ranks'))
         else:
             op_item['dtype'] = str(type(op_data.get('value')))
             op_item['shape'] = '[]'
@@ -166,7 +187,7 @@ def gen_op_item(op_data, op_name):
                 op_item[i] = op_data.get('value')
     if not op_item.get('md5'):
         op_item['md5'] = f"{zlib.crc32(str(op_data.get('value', '')).encode()):08x}"
-    
+
     return op_item
 
 
@@ -276,6 +297,22 @@ def result_item_init(n_info, b_info, dump_mode):
     return result_item
 
 
+def count_struct(op_dict):
+    parts = [
+        CompareConst.OP_NAME,
+        CompareConst.INPUT_STRUCT,
+        CompareConst.OUTPUT_STRUCT,
+        CompareConst.PARAMS_STRUCT,
+        CompareConst.PARAMS_GRAD_STRUCT
+    ]
+    lengths = [len(op_dict.get(part, [])) for part in parts]
+    num = lengths[0]
+    if num != sum(lengths[1:]):
+        logger.error(f"Length of names and structs of op_dict not match. Please check! op_dict: {op_dict}")
+        raise CompareException(CompareException.NAMES_STRUCTS_MATCH_ERROR)
+    return tuple(lengths)
+
+
 def get_accuracy(result, n_dict, b_dict, dump_mode):
     def get_accuracy_core(n_start, n_len, b_start, b_len, key):
         min_len = min(n_len, b_len)
@@ -355,31 +392,50 @@ def get_accuracy(result, n_dict, b_dict, dump_mode):
 
                 result.append(result_item)
 
-    n_num = len(n_dict['op_name'])
-    b_num = len(b_dict['op_name'])
-    n_num_input = len([name for name in n_dict['op_name']
-                       if Const.INPUT in name.split(Const.SEP) or Const.KWARGS in name.split(Const.SEP)])
-    b_num_input = len([name for name in b_dict['op_name']
-                       if Const.INPUT in name.split(Const.SEP) or Const.KWARGS in name.split(Const.SEP)])
-    n_num_output = n_num - n_num_input
-    b_num_output = b_num - b_num_input
-    get_accuracy_core(0, n_num_input, 0, b_num_input, 'input_struct')
-    get_accuracy_core(n_num_input, n_num_output, b_num_input, b_num_output, 'output_struct')
+    n_num, n_num_input, n_num_output, n_num_params, n_num_params_grad = count_struct(n_dict)
+    b_num, b_num_input, b_num_output, b_num_params, b_num_params_grad = count_struct(b_dict)
+
+    get_accuracy_core(0, n_num_input, 0, b_num_input, CompareConst.INPUT_STRUCT)
+    get_accuracy_core(n_num_input + n_num_output, n_num_params, b_num_input + b_num_output, b_num_params,
+                      CompareConst.PARAMS_STRUCT)
+    get_accuracy_core(n_num_input, n_num_output, b_num_input, b_num_output, CompareConst.OUTPUT_STRUCT)
+    get_accuracy_core(n_num_input + n_num_output + n_num_params, n_num_params_grad,
+                      b_num_input + b_num_output + b_num_params, b_num_params_grad,
+                      CompareConst.PARAMS_GRAD_STRUCT)
+
+
+def append_stack_info(result_item, npu_stack_info, index):
+    """添加堆栈信息到 result_item"""
+    if npu_stack_info and index == 0:
+        result_item.extend(npu_stack_info)
+    else:
+        result_item.append(CompareConst.NONE)
 
 
 def get_un_match_accuracy(result, n_dict, dump_mode):
-    index_out = 0
     npu_stack_info = n_dict.get("stack_info", None)
     bench_name, bench_type, bench_shape = CompareConst.N_A, CompareConst.N_A, CompareConst.N_A
-    err_msg = CompareConst.NO_BENCH
-    accuracy_check_res = CompareConst.N_A
-    for index, n_name in enumerate(n_dict["op_name"]):
-        name_ele_list = n_name.split(Const.SEP)
-        if Const.INPUT in name_ele_list or Const.KWARGS in name_ele_list:
-            n_struct = safe_get_value(n_dict, index, "n_dict", key=CompareConst.INPUT_STRUCT)
-        if Const.OUTPUT in name_ele_list:
-            n_struct = safe_get_value(n_dict, index_out, "n_dict", key=CompareConst.OUTPUT_STRUCT)
-            index_out += 1
+
+    struct_to_index_mapping = {
+        CompareConst.INPUT_STRUCT: 0,
+        CompareConst.OUTPUT_STRUCT: 0,
+        CompareConst.PARAMS_STRUCT: 0,
+        CompareConst.PARAMS_GRAD_STRUCT: 0
+    }
+
+    op_name_list = n_dict.get(CompareConst.OP_NAME)
+    summary_list = n_dict.get(Const.SUMMARY)
+    data_name_list = n_dict.get('data_name')
+    op_name_reorder, summary_reorder, _ = reorder_op_x_list(op_name_list,
+                                                            summary_list,
+                                                            data_name_list)
+    for index, n_name in enumerate(op_name_reorder):
+        _, state = get_name_and_state(n_name)
+        struct_key = CompareConst.STATE_TO_STRUCT_MAPPING.get(state)
+        if not struct_key:
+            continue
+        n_struct = safe_get_value(n_dict, struct_to_index_mapping.get(struct_key), "n_dict", key=struct_key)
+        struct_to_index_mapping[struct_key] += 1
 
         try:
             result_item = [n_name, bench_name, n_struct[0], bench_type, n_struct[1], bench_shape]
@@ -390,28 +446,26 @@ def get_un_match_accuracy(result, n_dict, dump_mode):
                       f"output_struct of n_dict is {n_dict[CompareConst.OUTPUT_STRUCT]}"
             logger.error(err_msg)
             raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR) from e
+
         if dump_mode == Const.MD5:
             result_item.extend([CompareConst.N_A] * 3)
-            if npu_stack_info and index == 0:
-                result_item.extend(npu_stack_info)
-            else:
-                result_item.append(CompareConst.NONE)
+            append_stack_info(result_item, npu_stack_info, index)
             result.append(result_item)
             continue
         if dump_mode == Const.SUMMARY:
             result_item.extend([CompareConst.N_A] * 8)
-        else:
+        if dump_mode == Const.ALL:
             result_item.extend([CompareConst.N_A] * 5)
-        npu_summary_data = safe_get_value(n_dict, index, "n_dict", key=CompareConst.SUMMARY)
-        result_item.extend(npu_summary_data)
+
+        npu_summary_data = safe_get_value(summary_reorder, index, "summary_reorder")
         bench_summary_data = [CompareConst.N_A] * 4
+        result_item.extend(npu_summary_data)
         result_item.extend(bench_summary_data)
+        err_msg = CompareConst.NO_BENCH
+        accuracy_check_res = CompareConst.N_A
         result_item.append(accuracy_check_res)
         result_item.append(err_msg)
-        if npu_stack_info and index == 0:
-            result_item.extend(npu_stack_info)
-        else:
-            result_item.append(CompareConst.NONE)
+        append_stack_info(result_item, npu_stack_info, index)
         if dump_mode == Const.ALL and result_item[1] == CompareConst.N_A:
             result_item.extend(["-1"])
         result.append(result_item)
@@ -423,6 +477,8 @@ def merge_tensor(tensor_list, dump_mode):
     op_dict[CompareConst.INPUT_STRUCT] = []
     op_dict[CompareConst.KWARGS_STRUCT] = []
     op_dict[CompareConst.OUTPUT_STRUCT] = []
+    op_dict[CompareConst.PARAMS_STRUCT] = []
+    op_dict[CompareConst.PARAMS_GRAD_STRUCT] = []
     op_dict[Const.SUMMARY] = []
     op_dict["stack_info"] = []
 
@@ -430,30 +486,25 @@ def merge_tensor(tensor_list, dump_mode):
         op_dict["data_name"] = []
 
     for tensor in tensor_list:
+        # A dict(len=2) with 'full_op_name' and 'full_info' is added to the tensor only if self.stack_mode is True
         if len(tensor) == 2:
             op_dict['stack_info'].append(tensor['full_info'])
             break
+
         op_dict["op_name"].append(tensor['full_op_name'])
-        name_ele_list = tensor['full_op_name'].split(Const.SEP)
-        name_to_struct_mapping = {
-            Const.INPUT: CompareConst.INPUT_STRUCT,
-            Const.KWARGS: CompareConst.KWARGS_STRUCT,
-            Const.OUTPUT: CompareConst.OUTPUT_STRUCT
-        }
-        for name_key, struct_key in name_to_struct_mapping.items():
-            if name_key in name_ele_list:
-                if dump_mode == Const.MD5:
-                    op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE], tensor[Const.MD5]))
-                else:
-                    op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE]))
-                break
+
+        _, state = get_name_and_state(tensor['full_op_name'])
+        struct_key = CompareConst.STATE_TO_STRUCT_MAPPING.get(state)
+        if not struct_key:
+            continue
+        if dump_mode == Const.MD5:
+            op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE], tensor[Const.MD5]))
+        else:
+            op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE]))
         op_dict[Const.SUMMARY].append([tensor[Const.MAX], tensor[Const.MIN], tensor[Const.MEAN], tensor[Const.NORM]])
 
         if dump_mode == Const.ALL:
             op_dict["data_name"].append(tensor['data_name'])
-            data_name = safe_get_value(op_dict, -1, "op_dict", key="data_name").rsplit(Const.SEP, 1)[0]
-            if data_name != "-1":
-                op_dict["op_name"][-1] = data_name
 
     if not op_dict[CompareConst.KWARGS_STRUCT]:
         del op_dict[CompareConst.KWARGS_STRUCT]
@@ -467,11 +518,90 @@ def print_compare_ends_info():
     logger.info('*' * total_len)
 
 
+def table_value_is_valid(value: str) -> bool:
+    if not isinstance(value, str):
+        return True
+    try:
+        # -1.00 or +1.00 should be consdiered as digit numbers
+        float(value)
+    except ValueError:
+        # otherwise, they will be considered as formular injections
+        return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
+    return True
+
+
+def get_name_and_state(name):
+    """
+    Get api/module name and state
+    example:
+    name = 'conv2d.forward.1.input.0'
+    return: ('conv2d.forward.1.', 'input')
+
+    name = 'Functional.pad.0.backward.output.0'
+    return: ('Functional.pad.0.backward.', 'output')
+
+    state type: input, output, kwargs, parameters, parameters_grad
+    """
+    if Const.PARAMS_GRAD in name.split(Const.SEP):
+        return name.split(Const.PARAMS_GRAD)[0], Const.PARAMS_GRAD
+
+    split = re.split(Const.REGEX_FORWARD_BACKWARD, name)
+    api = f'{split[0]}.{split[1]}.'
+    state_str = split[2]
+    match = re.match(r'^(\d+\.)?(input|output|kwargs|parameters)\..+$', state_str)
+    if not match:
+        raise CompareException(f'Invalid name string: {name}')
+    if match.group(1):
+        api = f'{api}{match.group(1)}'
+    state = match.group(2)
+    return api, state
+
+
+def reorder_op_name_list(op_name_list):
+    if not op_name_list:
+        return op_name_list
+
+    parameters = []
+    output = []
+    parameters_grad = []
+    others = []
+    for x in op_name_list:
+        state = get_name_and_state(x)[1]
+        if state == Const.PARAMS:
+            parameters.append(x)
+        elif state == Const.OUTPUT:
+            output.append(x)
+        elif state == Const.PARAMS_GRAD:
+            parameters_grad.append(x)
+        else:
+            others.append(x)
+    # 合并others, parameters, 和output，确保parameters排在output前面
+    op_name_reorder = others + parameters + output + parameters_grad
+    return op_name_reorder
+
+
+def reorder_op_x_list(op_name_list, summary_list, data_name_list):
+    """对op_name, summary, data_name重新排序，把parameters放到input后output前，data_name由于统计量比对时，为None，单独处理"""
+    if not op_name_list or not summary_list:
+        return op_name_list, summary_list, data_name_list
+
+    index_map = {name: index for index, name in enumerate(op_name_list)}
+
+    op_name_reorder = reorder_op_name_list(op_name_list)
+    summary_reorder = [summary_list[index_map.get(name)] for name in op_name_reorder]
+    if data_name_list:
+        data_name_reorder = [data_name_list[index_map.get(name)] for name in op_name_reorder]
+    else:
+        data_name_reorder = data_name_list
+
+    return op_name_reorder, summary_reorder, data_name_reorder
+
+
 def _compare_parser(parser):
     parser.add_argument("-i", "--input_path", dest="input_path", type=str,
                         help="<Required> The compare input path, a dict json.", required=True)
     parser.add_argument("-o", "--output_path", dest="output_path", type=str,
-                        help="<Required> The compare task result out path. Default path: ./output", 
+                        help="<Required> The compare task result out path. Default path: ./output",
                         required=False, default="./output", nargs="?", const="./output")
     parser.add_argument("-s", "--stack_mode", dest="stack_mode", action="store_true",
                         help="<optional> Whether to save stack info.", required=False)

msprobe/core/data_dump/data_collector.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");
@@ -38,6 +38,9 @@ class DataCollector:
         self.module_processor = DataProcessorFactory.get_module_processor(self.config.framework)
         self.module_count = {}
         self.scope = ScopeFactory(self.config).build_scope()
+        self.backward_module_names = {}
+        self.optimizer_status = ""
+        self.optimizer_status_first_start = {Const.OPTIMIZER: True, Const.CLIP_GRAD: True}
         atexit.register(self.write_json)
 
     @property
@@ -53,8 +56,15 @@ class DataCollector:
         return (not scope or scope.check(name)) and pid == os.getpid()
 
     @staticmethod
-    def is_inplace(module):
-        return getattr(module, "op_is_inplace", False)
+    def set_is_recomputable(data_info, is_recompute):
+        if data_info and len(data_info) == 1 and is_recompute is not None: # 正常情况下data_info的长度应改为1
+            data_info[list(data_info.keys())[0]]["is_recompute"] = is_recompute
+
+    def reset_status(self):
+        self.optimizer_status = ""
+        self.optimizer_status_first_start = {Const.OPTIMIZER: True, Const.CLIP_GRAD: True}
+        self.data_writer.reset_cache()
+        self.backward_module_names.clear()
 
     def if_return_forward_new_output(self):
         return self.data_processor.if_return_forward_new_output()
@@ -79,69 +89,105 @@ class DataCollector:
         logger.debug(msg)
         self.data_writer.update_data(data_info)
 
-    def pre_forward_data_collect(self, name, module, pid, module_input_output):
-        if self.config.level == Const.LEVEL_L2 and self.check_scope_and_pid(self.scope, name, pid):
-            self.data_processor.analyze_pre_forward(name, module, module_input_output)
+    def forward_input_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
+        if self.config.task == Const.FREE_BENCHMARK:
+            backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
+            if self.check_scope_and_pid(self.scope, backward_name, pid):
+                self.data_processor.analyze_forward_input(backward_name, module, module_input_output)
+            return
+
+        if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
-        if self.check_scope_and_pid(self.scope, backward_name, pid):
-            self.data_processor.analyze_pre_forward(backward_name, module, module_input_output)
-        if not self.is_inplace(module) or not self.check_scope_and_pid(self.scope, name, pid):
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
+            data_info = self.data_processor.analyze_forward_input(name, module, module_input_output)
+        self.set_is_recomputable(data_info, is_recompute)
+        if self.config.level == Const.LEVEL_L2:
             return
-        logger.info(f"API {name} is inplace.")
-        data_info = self.data_processor.analyze_pre_forward_inplace(name, module_input_output)
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
-    def forward_data_collect(self, name, module, pid, module_input_output):
+    def forward_output_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
+
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
+            data_info = self.data_processor.analyze_forward_output(name, module, module_input_output)
+        self.set_is_recomputable(data_info, is_recompute)
         if self.config.level == Const.LEVEL_L2:
-            self.data_processor.analyze_forward(name, module, module_input_output)
+            return
+        self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
+        self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
+
+    def forward_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
+        self.update_construct(name)
+        if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        if not self.is_inplace(module):
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
             data_info = self.data_processor.analyze_forward(name, module, module_input_output)
-        else:
-            data_info = self.data_processor.analyze_forward_inplace(name, module_input_output)
+        self.set_is_recomputable(data_info, is_recompute)
         self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
-    def backward_data_collect(self, name, module, pid, module_input_output):
+    def backward_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        data_info = self.data_processor.analyze_backward(name, module, module_input_output)
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
+            data_info = self.data_processor.analyze_backward(name, module, module_input_output)
         if self.config.level == Const.LEVEL_L2:
             return
+        # 获取执行反向的模块名称
+        if data_info and name.split(Const.SEP)[0] in Const.MODULE_PREFIX:
+            module_name = name.rsplit(Const.SEP, 2)[0]
+            # 将模块名称加入到反向模块名称集合中，用于梯度收集时判断是否需要收集梯度
+            self.backward_module_names[module_name] = True
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
-    def backward_input_data_collect(self, name, module, pid, module_input_output):
+    def backward_input_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        data_info = self.data_processor.analyze_backward_input(name, module, module_input_output)
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
+            data_info = self.data_processor.analyze_backward_input(name, module, module_input_output)
+        self.set_is_recomputable(data_info, is_recompute)
         self.handle_data(name, data_info)
 
-    def backward_output_data_collect(self, name, module, pid, module_input_output):
+    def backward_output_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        data_info = self.data_processor.analyze_backward_output(name, module, module_input_output)
+        data_info = {}
+        if self.config.task != Const.STRUCTURE:
+            data_info = self.data_processor.analyze_backward_output(name, module, module_input_output)
+        self.set_is_recomputable(data_info, is_recompute)
         self.handle_data(name, data_info)
 
     def update_construct(self, name):
         if self.config.level not in DataCollector.level_without_construct:
-            self.data_writer.update_construct({name: self.module_processor.api_parent_node})
+            if self.optimizer_status in [Const.OPTIMIZER, Const.CLIP_GRAD]:
+                if self.optimizer_status_first_start[self.optimizer_status]:
+                    self.data_writer.update_construct({self.optimizer_status: None})
+                    self.optimizer_status_first_start[self.optimizer_status] = False
+                self.data_writer.update_construct({name: self.optimizer_status})
+            else:
+                self.data_writer.update_construct({name: self.module_processor.api_parent_node})
             self.data_writer.update_construct(self.module_processor.module_node)
 
     def handle_data(self, name, data_info, flush=False):
         if data_info:
             self.update_data(name, data_info)
+        if self.config.async_dump:
+            return
         if not flush:
             self.data_writer.flush_data_periodically()
         else:
@@ -149,7 +195,36 @@ class DataCollector:
 
     def update_dump_paths(self, *args):
         self.data_writer.update_dump_paths(*args)
-        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level)
+
+    def initialize_json_file(self, framework=Const.UNKNOWN_FRAMEWORK):
+        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level, framework=framework)
 
     def update_iter(self, current_iter):
         self.data_processor.update_iter(current_iter)
+
+    def params_data_collect(self, name, param_name, pid, data):
+        grad_name = name + Const.SEP + Const.PARAMS_GRAD
+        # 校验scope和pid，以及当前name是否有过反向计算
+        if not self.check_scope_and_pid(self.scope, name, pid) and not self.backward_module_names.get(name):
+            # 如果没有反向计算，则需要清除之前占位写入的grad数据
+            if self.data_writer.cache_data.get("data"):
+                self.data_writer.cache_data.get("data").pop(grad_name, None)
+            return
+        data_info = self.data_processor.analyze_params(grad_name, param_name, data)
+        self.handle_data(grad_name, data_info, flush=self.data_processor.is_terminated)
+
+    def fill_stack_tensor_data(self):
+        self.data_writer.fill_stack_tensor_data()
+
+    def debug_data_collect_forward(self, variable, name_with_count):
+
+        data_info = self.data_processor.analyze_debug_forward(variable, name_with_count)
+        self.data_writer.update_debug({name_with_count: data_info})
+
+    def debug_data_collect_backward(self, variable, grad_name_with_count):
+        # prepare all None nested data structure
+        all_none_data_info = self.data_processor.analyze_element_to_all_none(variable)
+        self.data_writer.update_debug({grad_name_with_count: all_none_data_info})
+
+        # register tensor backward hook
+        self.data_processor.analyze_debug_backward(variable, grad_name_with_count, self.data_writer.cache_debug['data'])