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

msprobe/pytorch/bench_functions/rotary_mul.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 
 
@@ -25,15 +40,22 @@ def npu_rotary_mul_backward(dy_tensor, x, r1, r2):
     x_shape = x.shape
     h = x.float()
     grad = dy_tensor.float()
-    condition_1 = (((r1_shape[0] == 1 and x_shape[0] != 1) or (r1_shape[0] == 1 and x_shape[0] == 1)) and
-                   ((r1_shape[2] == 1 and x_shape[2] != 1) or (r1_shape[2] == 1 and x_shape[2] == 1)) and
-                   (r1_shape[1] == x_shape[1]) and (r1_shape[3] == x_shape[3]))
-    condition_2 = (((r1_shape[0] == 1 and x_shape[0] != 1) or (r1_shape[0] == 1 and x_shape[0] == 1)) and
-                   ((r1_shape[1] == 1 and x_shape[1] != 1) or (r1_shape[1] == 1 and x_shape[1] == 1)) and
-                   (r1_shape[2] == x_shape[2]) and (r1_shape[3] == x_shape[3]))
-    condition_3 = (((r1_shape[2] == 1 and x_shape[2] != 1) or (r1_shape[2] == 1 and x_shape[2] == 1)) and
-                   ((r1_shape[1] == 1 and x_shape[1] != 1) or (r1_shape[1] == 1 and x_shape[1] == 1)) and
-                   (r1_shape[0] == x_shape[0]) and (r1_shape[3] == x_shape[3]))
+    if len(r1_shape) < 4 or len(x_shape) < 4:
+        raise RuntimeError(f"Shape of r1 and x should at least be 4-dimension, "
+                           f"but got r1 shape:{r1_shape}, x shape:{x_shape}")
+    condition_1 = (r1_shape[0] == 1
+                   and r1_shape[1] == x_shape[1]
+                   and r1_shape[2] == 1
+                   and r1_shape[3] == x_shape[3])
+    condition_2 = (r1_shape[0] == 1
+                   and r1_shape[1] == 1
+                   and r1_shape[2] == x_shape[2]
+                   and r1_shape[3] == x_shape[3])
+    condition_3 = (r1_shape[0] == x_shape[0]
+                   and r1_shape[1] == 1
+                   and r1_shape[2] == 1
+                   and r1_shape[3] == x_shape[3])
+
     if condition_1:
         for i in range(x_shape[0]):
             for j in range(x_shape[2]):
@@ -49,4 +71,5 @@ def npu_rotary_mul_backward(dy_tensor, x, r1, r2):
             for j in range(x_shape[2]):
                 r2_grad[:, 0, 0, :] += (x_new2[:, i, j, :] * grad[:, i, j, :])
                 r1_grad[:, 0, 0, :] += (h[:, i, j, :] * grad[:, i, j, :])
+
     return x.grad.cpu(), r1_grad.cpu(), r2_grad.cpu()

msprobe/pytorch/bench_functions/scaled_mask_softmax.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 
 

msprobe/pytorch/bench_functions/swiglu.py

@@ -1,16 +1,35 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 
 
 def npu_swiglu(x, dim=-1):
     tensor_dtype = x.dtype
 
-    inTensors = torch.chunk(x, 2, dim=dim)
+    try:
+        in_tensors = torch.chunk(x, 2, dim=dim)
+    except Exception as e:
+        raise RuntimeError(f"Invalid chunk x into 2 tensors with shape {x.shape} and dimension {dim}") from e
+
     if tensor_dtype == torch.float32:
-        tensor_scalar = torch.sigmoid(torch.mul(inTensors[0], 1.0))
-        output_data = torch.mul(torch.mul(tensor_scalar, inTensors[0]), inTensors[1])
+        tensor_scalar = torch.sigmoid(torch.mul(in_tensors[0], 1.0))
+        output_data = torch.mul(torch.mul(tensor_scalar, in_tensors[0]), in_tensors[1])
     else:
-        tensor_self_float = inTensors[0].type(torch.float)
-        tensor_other_float = inTensors[1].type(torch.float)
+        tensor_self_float = in_tensors[0].type(torch.float)
+        tensor_other_float = in_tensors[1].type(torch.float)
         tensor_out_float = torch.nn.functional.silu(tensor_self_float).type(tensor_dtype).type(
             torch.float32) * tensor_other_float
         output_data = tensor_out_float.type(tensor_dtype)
@@ -19,7 +38,11 @@ def npu_swiglu(x, dim=-1):
 
 def npu_swiglu_backward(grad, x, dim=-1):
     tensor_dtype = grad.dtype
-    in_tensors = torch.chunk(x, 2, dim=dim)
+    try:
+        in_tensors = torch.chunk(x, 2, dim=dim)
+    except Exception as e:
+        raise RuntimeError(f"Invalid chunk x into 2 tensors with shape {x.shape} and dimension {dim}") from e
+
     tensor_grad_out = grad
 
     if tensor_dtype == torch.float16:

msprobe/pytorch/common/__init__.py

@@ -1,2 +1,17 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from .parse_json import parse_json_info_forward_backward
 from .utils import seed_all

msprobe/pytorch/common/log.py

@@ -1,9 +1,21 @@
-import os
-import time
-import sys
-from msprobe.pytorch.common.utils import get_rank_if_initialized
-from msprobe.core.common.log import BaseLogger
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from msprobe.core.common.exceptions import DistributedNotInitializedError
+from msprobe.core.common.log import BaseLogger
+from msprobe.pytorch.common.utils import get_rank_if_initialized
 
 
 class PyTorchLogger(BaseLogger):
@@ -18,4 +30,4 @@ class PyTorchLogger(BaseLogger):
         return current_rank
 
 
-logger = PyTorchLogger()
+logger = PyTorchLogger()

msprobe/pytorch/common/parse_json.py

@@ -1,25 +1,32 @@
-import json
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 from msprobe.core.common.exceptions import ParseJsonException
-from msprobe.core.common.file_utils import FileOpen
+from msprobe.core.common.file_utils import load_json
+from msprobe.core.common.log import logger
 
 
 def parse_json_info_forward_backward(json_path):
-    def parse_data_name_with_pattern(data_name, pattern):
-        name_struct = data_name.split('.')
-        if not name_struct[-1] == pattern:
-            raise ParseJsonException(ParseJsonException.UnexpectedNameStruct,
-                f"{data_name} in file {json_path}")
-        api_name = '.'.join(name_struct[:-1])
-        return api_name
-
-    with FileOpen(json_path, 'r') as f:
-        dump_json = json.load(f)
+    dump_json = load_json(json_path)
 
     real_data_path = dump_json.get("dump_data_dir")
     dump_data = dump_json.get("data")
+    if dump_data is None:
+        raise ParseJsonException(ParseJsonException.InvalidDumpJson, "something wrong with dump, no data found in dump.json")
     if not dump_data:
-        raise ParseJsonException(ParseJsonException.InvalidDumpJson, "dump数据中没有data字段")
+        logger.warning("data field is empty, no overflow data found.")
 
     forward_data = {}
     backward_data = {}
@@ -27,13 +34,21 @@ def parse_json_info_forward_backward(json_path):
         if "Module" in data_name:
             continue
         if "forward" in data_name:
-            api_name = parse_data_name_with_pattern(data_name, "forward")
+            api_name = parse_data_name_with_pattern(data_name, "forward", json_path)
             forward_data.update({api_name: data_item})
         elif "backward" in data_name:
-            api_name = parse_data_name_with_pattern(data_name, "backward")
+            api_name = parse_data_name_with_pattern(data_name, "backward", json_path)
             backward_data.update({api_name: data_item})
         else:
             raise ParseJsonException(ParseJsonException.UnexpectedNameStruct,
-                f"{data_name} in file {json_path}.")
+                                     f"{data_name} in file {json_path}.")
 
     return forward_data, backward_data, real_data_path
+
+
+def parse_data_name_with_pattern(data_name, pattern, json_path):
+    name_struct = data_name.split('.')
+    if not name_struct[-1] == pattern:
+        raise ParseJsonException(ParseJsonException.UnexpectedNameStruct, f"{data_name} in file {json_path}")
+    api_name = '.'.join(name_struct[:-1])
+    return api_name

msprobe/pytorch/common/utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,20 +12,23 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
+
 import io
 import os
+import pickle
 import random
 import stat
+from functools import wraps
+
+import numpy as np
 import torch
 import torch.distributed as dist
-import numpy as np
-from functools import wraps
 from msprobe.core.common.exceptions import DistributedNotInitializedError
-from msprobe.core.common.log import logger
 from msprobe.core.common.file_utils import (FileCheckConst, change_mode,
-                                            check_file_or_directory_path, check_path_before_create)
-
+                                            check_file_or_directory_path, check_path_before_create, FileOpen)
+from msprobe.core.common.log import logger
+from msprobe.core.common.utils import check_seed_all
+from packaging import version
 
 try:
     import torch_npu
@@ -35,10 +37,8 @@ except ImportError:
 else:
     is_gpu = False
 
-
 torch_without_guard_version = torch.__version__ >= '2.1'
 
-
 if not is_gpu and not torch_without_guard_version:
     from torch_npu.utils.device_guard import torch_device_guard as torch_npu_device_guard
 
@@ -46,7 +46,6 @@ npu_distributed_api = ['isend', 'irecv']
 
 
 def parameter_adapter(func):
-
     def handle_masked_select(input_tensor, indices):
         masked_select_func = getattr(torch._C._VariableFunctionsClass, "masked_select")
         if input_tensor.dtype == torch.bfloat16:
@@ -77,20 +76,22 @@ def parameter_adapter(func):
                 else:
                     res = [input_tensor[tensor_index] for tensor_index in indices]
                     return getattr(torch._C._VariableFunctionsClass, "stack")(res, 0)
-        if self.op_name_ == "__eq__" and args[1] is None:
+        if self.op_name_ == "__eq__" and len(args) > 1 and args[1] is None:
             return False
         return func(self, *args, **kwargs)
+
     return inner
 
 
 def torch_device_guard(func):
     if is_gpu or torch_without_guard_version:
         return func
-    # Parse args/kwargs matched torch.device objects
 
+    # Parse args/kwargs matched torch.device objects
     @torch_npu_device_guard
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
+
     return wrapper
 
 
@@ -105,20 +106,28 @@ def get_rank_if_initialized():
 
 
 def seed_all(seed=1234, mode=False):
-    random.seed(seed)
-    os.environ['PYTHONHASHSEED'] = str(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.use_deterministic_algorithms(mode)
-    if is_gpu:
-        torch.cuda.manual_seed_all(seed)
-        torch.cuda.manual_seed(seed)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.enable = False
-        torch.backends.cudnn.benchmark = False
-    else:
-        torch_npu.npu.manual_seed_all(seed)
-        torch_npu.npu.manual_seed(seed)
+    check_seed_all(seed, mode)
+    try:
+        random.seed(seed)
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        cuda_version = torch.version.cuda
+        if cuda_version is not None and version.parse(cuda_version) >= version.parse("10.2"):
+            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+        os.environ['HCCL_DETERMINISTIC'] = str(mode)
+        torch.use_deterministic_algorithms(mode)
+        if is_gpu:
+            torch.cuda.manual_seed_all(seed)
+            torch.cuda.manual_seed(seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.enable = False
+            torch.backends.cudnn.benchmark = False
+        else:
+            torch_npu.npu.manual_seed_all(seed)
+            torch_npu.npu.manual_seed(seed)
+    except Exception as e:
+        logger.error(f"There is an unexpected error while determinating randomness. {e}")
 
 
 class Const:
@@ -191,10 +200,7 @@ class Const:
     ENV_ENABLE = "1"
     ENV_DISABLE = "0"
 
-    MAX_SEED_VALUE = 2**32 - 1
-
-    INPLACE_LIST = ["broadcast", "all_reduce", "reduce", "all_gather", "gather", "scatter", "reduce_scatter",
-                    "_reduce_scatter_base", "_all_gather_base", "all_to_all_single"]
+    MAX_SEED_VALUE = 2 ** 32 - 1
 
     TASK_LIST = ["tensor", "statistics", "overflow_check", "free_benchmark"]
     LEVEL_LIST = ["L0", "L1", "L2", "mix"]
@@ -257,34 +263,84 @@ def print_rank_0(message):
             logger.info(message)
     else:
         logger.info(message)
-        
+
 
 def load_pt(pt_path, to_cpu=False):
     pt_path = os.path.realpath(pt_path)
     check_file_or_directory_path(pt_path)
     try:
         if to_cpu:
-            pt = torch.load(pt_path, map_location=torch.device("cpu"))
+            pt = torch.load(pt_path, map_location=torch.device("cpu"), weights_only=True)
         else:
-            pt = torch.load(pt_path)
+            pt = torch.load(pt_path, weights_only=True)
     except Exception as e:
         raise RuntimeError(f"load pt file {pt_path} failed") from e
     return pt
 
 
 def save_pt(tensor, filepath):
-    filepath = os.path.realpath(filepath)
     check_path_before_create(filepath)
+    filepath = os.path.realpath(filepath)
     try:
         torch.save(tensor, filepath)
     except Exception as e:
         logger.error("Save pt file failed, please check according possible error causes: "
-                            "1. out of disk space or disk error, "
-                            "2. no permission to write files, etc.")
+                     "1. out of disk space or disk error, "
+                     "2. no permission to write files, etc.")
         raise RuntimeError(f"save pt file {filepath} failed") from e
     change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
+class TypeCheckingUnpickler(pickle.Unpickler):
+    """
+    This class is a subclass of pickle.Unpickler, which is used to unpickle pickled objects.
+    It overrides the find_class method to add type checking functionality.
+    """
+    allowed_types = [
+        "str",
+        "ApiData",
+        "OrderedDict",
+        "_rebuild_tensor_v2",  # from torch.utils
+        "_load_from_bytes"  # from torch.storage
+    ]
+
+    def find_class(self, module, name):
+        """
+        Method to find the class of the object to be unpickled.
+        Throws pickle.UnpicklingError If the object type is not in the allowed types list.
+        """
+        if name in self.allowed_types:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError("Unsupported object type: {}.{}".format(module, name))
+
+
+def save_pkl(tensor, filepath):
+    """Save ApiData or str objection by pickle"""
+    check_path_before_create(filepath)
+    filepath = os.path.realpath(filepath)
+    try:
+        with FileOpen(filepath, 'wb') as f:
+            pickle.dump(tensor, f)
+    except Exception as e:
+        logger.error("Save pt file failed, please check according possible error causes: "
+                     "1. out of disk space or disk error, "
+                     "2. no permission to write files, etc.")
+        raise RuntimeError(f"save pt file {filepath} failed") from e
+    change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
+def load_pkl(pt_path):
+    """Load ApiData or str objection by pickle for accuracy_checker_online"""
+    check_file_or_directory_path(pt_path)
+    pt_path = os.path.realpath(pt_path)
+    try:
+        with FileOpen(pt_path, 'rb') as f:
+            pt = TypeCheckingUnpickler(f).load()
+    except Exception as e:
+        raise RuntimeError(f"load pt file {pt_path} failed: {e}") from e
+    return pt
+
+
 def save_api_data(api_data):
     """Save data to io stream"""
     try:

msprobe/pytorch/compare/distributed_compare.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2019-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2019-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,14 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
+
 import os
 from msprobe.core.common.utils import CompareException, check_compare_param, \
-    check_configuration_param, task_dumppath_get
+    check_configuration_param, set_dump_path, get_dump_mode
 from msprobe.core.common.file_utils import create_directory
 from msprobe.core.common.exceptions import FileCheckException
 from msprobe.pytorch.common.log import logger
-from msprobe.core.common.const import Const
 from msprobe.pytorch.compare.pt_compare import PTComparator
 from msprobe.core.compare.utils import check_and_return_dir_contents, extract_json
 
@@ -32,6 +30,7 @@ def compare_distributed(npu_dump_dir, bench_dump_dir, output_path, **kwargs):
     stack_mode = kwargs.get('stack_mode', False)
     auto_analyze = kwargs.get('auto_analyze', True)
     fuzzy_match = kwargs.get('fuzzy_match', False)
+    is_print_compare_log = kwargs.get('is_print_compare_log', True)
     # get the ranks and match by order
     npu_ranks = sorted(check_and_return_dir_contents(npu_dump_dir, 'rank'))
     bench_ranks = sorted(check_and_return_dir_contents(bench_dump_dir, 'rank'))
@@ -51,16 +50,16 @@ def compare_distributed(npu_dump_dir, bench_dump_dir, output_path, **kwargs):
             'npu_json_path': npu_path,
             'bench_json_path': bench_path,
             'stack_json_path': stack_path,
-            'is_print_compare_log': True
+            'is_print_compare_log': is_print_compare_log
         }
         try:
-            summary_compare, md5_compare = task_dumppath_get(dump_result_param)
-            check_configuration_param(stack_mode, auto_analyze, fuzzy_match)
+            set_dump_path(dump_result_param)
+            dump_mode = get_dump_mode(dump_result_param)
+            check_configuration_param(stack_mode, auto_analyze, fuzzy_match, is_print_compare_log)
             create_directory(output_path)
-            check_compare_param(dump_result_param, output_path, summary_compare=summary_compare, md5_compare=md5_compare)
+            check_compare_param(dump_result_param, output_path, dump_mode)
         except (CompareException, FileCheckException) as error:
             logger.error('Compare failed. Please check the arguments and do it again!')
             raise CompareException(error.code) from error
         pt_comparator = PTComparator()
-        pt_comparator.compare_core(dump_result_param, output_path, suffix=f'_{nr}-{br}', summary_compare=summary_compare,
-                     md5_compare=md5_compare, **kwargs)
+        pt_comparator.compare_core(dump_result_param, output_path, suffix=f'_{nr}-{br}', dump_mode=dump_mode, **kwargs)

msprobe/pytorch/compare/match.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 from msprobe.core.common.utils import CompareException
 from msprobe.core.common.file_utils import load_yaml

msprobe/pytorch/compare/pt_compare.py

@@ -1,19 +1,52 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os.path
 import torch
 from msprobe.core.common.const import FileCheckConst
 from msprobe.pytorch.common.log import logger
 from msprobe.core.common.exceptions import FileCheckException
 from msprobe.core.compare.acc_compare import Comparator
-from msprobe.core.common.utils import check_configuration_param, task_dumppath_get, check_compare_param, CompareException
-from msprobe.core.common.file_utils import FileChecker, create_directory
+from msprobe.core.common.utils import check_configuration_param, check_compare_param, \
+    CompareException, set_dump_path, get_dump_mode
+from msprobe.core.common.file_utils import FileChecker, create_directory, load_yaml
 from msprobe.pytorch.common.utils import load_pt
 
 
 class PTComparator (Comparator):
-    def __init__(self):
+    def __init__(self, data_mapping=None):
         self.frame_name = PTComparator.__name__
+        self.data_mapping = data_mapping
+        if isinstance(self.data_mapping, str) or self.data_mapping is None:
+            self.data_mapping_dict = self.load_mapping_file(self.data_mapping)
+        elif isinstance(self.data_mapping, dict):
+            self.data_mapping_dict = self.data_mapping
+        else:
+            raise TypeError(f"The type of parameter `data_mapping` must be dict, str or None, but got "
+                            f"{type(self.data_mapping)}")
+
+    def load_mapping_file(self, mapping_file):
+        if isinstance(mapping_file, str):
+            mapping_dict = load_yaml(mapping_file)
+        else:
+            mapping_dict = {}
+        return mapping_dict
     
     def read_npy_data(self, dir_path, file_name):
+        if not file_name:
+            return None
         data_path = os.path.join(dir_path, file_name)
         path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
                                 FileCheckConst.PT_SUFFIX, False)
@@ -35,16 +68,17 @@ class PTComparator (Comparator):
         return data_value  
     
     
-def compare(input_param, output_path, stack_mode=False, auto_analyze=True, fuzzy_match=False):
+def compare(input_param, output_path, stack_mode=False, auto_analyze=True, fuzzy_match=False, **kwargs):
     try:
-        summary_compare, md5_compare = task_dumppath_get(input_param)
-        check_configuration_param(stack_mode, auto_analyze, fuzzy_match)
+        set_dump_path(input_param)
+        dump_mode = get_dump_mode(input_param)
+        check_configuration_param(stack_mode, auto_analyze, fuzzy_match, input_param.get('is_print_compare_log', True))
         create_directory(output_path)
-        check_compare_param(input_param, output_path, summary_compare, md5_compare)
+        check_compare_param(input_param, output_path, dump_mode)
+        data_mapping = kwargs.get('data_mapping', None)
     except (CompareException, FileCheckException) as error:
         logger.error('Compare failed. Please check the arguments and do it again!')
         raise CompareException(error.code) from error
-    pt_comparator = PTComparator()
+    pt_comparator = PTComparator(data_mapping)
     pt_comparator.compare_core(input_param, output_path, stack_mode=stack_mode,
-                 auto_analyze=auto_analyze, fuzzy_match=fuzzy_match, summary_compare=summary_compare,
-                 md5_compare=md5_compare)
+                 auto_analyze=auto_analyze, fuzzy_match=fuzzy_match, dump_mode=dump_mode)