mindstudio-probe 1.3.0__py3-none-any.whl → 8.1.1__py3-none-any.whl

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/device_dispatch.py

@@ -29,7 +29,6 @@ from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import move2target_device
 from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import generate_cpu_params
 
-
 # NPU vs GPU api list
 CompareApi = set(absolute_standard_api) | set(binary_standard_api) | set(thousandth_standard_api)
 
@@ -43,6 +42,15 @@ OnlineApiPrecisionCompareConfig = namedtuple('OnlineApiPrecisionCompareConfig',
 CommonCompareConfig = namedtuple('CommonCompareConfig', ['compare', 'handle_func', 'config'])
 
 
+def get_gpu_device():
+    try:
+        import torch_npu
+        is_gpu = False
+    except ImportError:
+        is_gpu = True
+    return is_gpu
+
+
 def run_ut_process(xpu_id, consumer_queue, common_config, api_precision_csv_file):
     """ When consumer_queue(shared with ConsumerDispatcher) is not empty, consume api data from consumer_queue.
     :param xpu_id: int
@@ -51,7 +59,9 @@ def run_ut_process(xpu_id, consumer_queue, common_config, api_precision_csv_file
     :param api_precision_csv_file: list, length is 2, result file name and details file name
     :return:
     """
-    gpu_device = torch.device(f'cuda:{xpu_id}')
+    device_info = "cuda" if get_gpu_device() else "npu"
+    logger.info(f"Start run_ut_process for {device_info} device, rank: {xpu_id}.")
+    gpu_device = torch.device(f'{device_info}:{xpu_id}')
 
     while True:
         if consumer_queue.empty():

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/server.py

@@ -12,19 +12,19 @@
 # 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
+from functools import partial
+import os
 import struct
-import hashlib
+import zlib
 import time
 import io
 from threading import Thread
 
-from twisted.internet import reactor, protocol, endpoints
+from twisted.internet import reactor, protocol, endpoints, ssl
 
 from msprobe.pytorch.common.utils import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.utils import cipher_list, \
-    STRUCT_UNPACK_MODE as unpack_mode, STR_TO_BYTES_ORDER as bytes_order
+    STRUCT_UNPACK_MODE as unpack_mode, STR_TO_BYTES_ORDER as bytes_order, verify_callback, load_ssl_pem
 
 
 class TCPServer:
@@ -44,15 +44,28 @@ class TCPServer:
         self.factory.protocol = self.build_protocol
 
         if self.tls_path:
-            from OpenSSL import SSL
-            from twisted.internet import ssl
-            server_key = os.path.join(self.tls_path, "server.key")
-            server_crt = os.path.join(self.tls_path, "server.crt")
-            server_context_factory = ssl.DefaultOpenSSLContextFactory(server_key, server_crt, SSL.TLSv1_2_METHOD)
-            server_context_ = server_context_factory.getContext()
-            server_context_.set_cipher_list(cipher_list)
-            server_context_.set_options(SSL.OP_NO_RENEGOTIATION)
-            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, server_context_factory)
+            server_key, server_crt, ca_crt, crl_pem = load_ssl_pem(
+                key_file=os.path.join(self.tls_path, "server.key"),
+                cert_file=os.path.join(self.tls_path, "server.crt"),
+                ca_file=os.path.join(self.tls_path, "ca.crt"),
+                crl_file=os.path.join(self.tls_path, "crl.pem")
+            )
+
+            ssl_options = ssl.CertificateOptions(
+                privateKey=server_key,
+                certificate=server_crt,
+                method=ssl.SSL.TLSv1_2_METHOD,
+                verify=True,
+                requireCertificate=True,
+                caCerts=[ca_crt],  # 信任的CA证书列表
+            )
+            ssl_context = ssl_options.getContext()
+            ssl_context.set_cipher_list(cipher_list)
+            ssl_context.set_options(ssl.SSL.OP_NO_RENEGOTIATION)
+            ssl_context.set_verify(ssl.SSL.VERIFY_PEER | ssl.SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
+                                   partial(verify_callback, crl=crl_pem))
+
+            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, ssl_options)
         else:
             endpoint = endpoints.TCP4ServerEndpoint(reactor, self.port)
         endpoint.listen(self.factory)
@@ -85,10 +98,10 @@ class ServerProtocol(protocol.Protocol):
         self.consumer_queue = shared_queue
         self.check_sum = check_sum
         self.length_width = 8
-        self.md5_width = 32
+        self.crc_width = 8
         self.obj_length = None
         self.tell = 0
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
         self.sequence_number = -1
         self.rank = -1
@@ -99,7 +112,7 @@ class ServerProtocol(protocol.Protocol):
         self.buffer = io.BytesIO()
         self.obj_length = None
         self.tell = 0
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
         self.factory.transport_dict[self.transport] = 1
         self.factory.transport_list.append(self.transport)
@@ -132,11 +145,12 @@ class ServerProtocol(protocol.Protocol):
             time.sleep(0.1)
 
         obj_key = str(self.sequence_number) + "_" + str(self.rank) + "_" + str(self.step)
+        # get the crc value of a 16-bit string with a length of 8
+        recv_crc = f"{zlib.crc32(self.obj_body):08x}"
 
-        recv_md5 = hashlib.md5(self.obj_body).hexdigest()
-        if self.check_sum and recv_md5 != self.obj_md5:
-            # when needs check md5 and check no pass, indicates received data error, send b"ERROR" to client.
-            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_md5}, but get {recv_md5}")
+        if self.check_sum and recv_crc != self.obj_crc:
+            # when needs check hash value and check no pass, indicates received data error, send b"ERROR" to client.
+            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_crc}, but get {recv_crc}")
             self.send_ack(self.ACK_ERROR)
         else:
             if self.obj_body == self.ACK_STOP:
@@ -146,7 +160,7 @@ class ServerProtocol(protocol.Protocol):
             if obj_key in self.sequence_number_dict:
                 logger.debug(f"这是一次异常的重传，可以忽略。 {obj_key}, {self.sequence_number_dict}")
             else:
-                self.sequence_number_dict[obj_key] = self.obj_md5
+                self.sequence_number_dict[obj_key] = self.obj_crc
                 self.consumer_queue.put(self.obj_body, block=True)
 
         self.reset_env()
@@ -173,7 +187,7 @@ class ServerProtocol(protocol.Protocol):
         self.sequence_number = -1
         self.rank = -1
         self.step = -1
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
 
     def dataReceived(self, data):
@@ -192,15 +206,15 @@ class ServerProtocol(protocol.Protocol):
             logger.debug(
                 f"流水号: {self.sequence_number}; RANK: {self.rank}; STEP: {self.step}; Length: {self.obj_length}")
 
-        # If needs check md5 but not parse md5 yet, read 32b md5 values
-        check_sum_and_md5 = (self.check_sum
+        # If needs check hash but not parse crc yet, read 8b crc values
+        check_sum_and_crc = (self.check_sum
                              and self.obj_length is not None
-                             and self.obj_md5 is None
-                             and len(self.buffer.getvalue()) - self.tell >= self.md5_width)
-        if check_sum_and_md5:
-            self.obj_md5 = self.buffer.read(self.md5_width).decode()
-            self.tell += self.md5_width
-            logger.debug(f"MD5: {self.obj_md5}")
+                             and self.obj_crc is None
+                             and len(self.buffer.getvalue()) - self.tell >= self.crc_width)
+        if check_sum_and_crc:
+            self.obj_crc = self.buffer.read(self.crc_width).decode()
+            self.tell += self.crc_width
+            logger.debug(f"Hash value: {self.obj_crc}")
 
         current_length = len(self.buffer.getvalue()) - self.tell
         if self.obj_length is not None and 0 < self.obj_length <= current_length:

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/utils.py

@@ -12,6 +12,17 @@
 # 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 gc
+import os
+from datetime import datetime, timezone
+
+from OpenSSL import crypto
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+from dateutil import parser
+
+from msprobe.core.common.file_utils import FileOpen
+from msprobe.core.common.log import logger
 
 cipher_list = ":".join(
     ["TLS_DHE_RSA_WITH_AES_128_GCM_SHA256",
@@ -42,3 +53,148 @@ cipher_list = ":".join(
 
 STRUCT_UNPACK_MODE = "!Q"
 STR_TO_BYTES_ORDER = "big"
+
+
+def is_certificate_revoked(cert, crl):
+    # 获取证书的序列号
+    cert_serial_number = cert.get_serial_number()
+
+    # 检查证书是否在CRL中
+    revoked_serials = [revoked_cert.serial_number for revoked_cert in crl]
+    if cert_serial_number in revoked_serials:
+        logger.error(f"证书已吊销:{cert_serial_number:020x}")
+        return True
+
+    return False
+
+
+def verify_callback(conn, cert, errno, depth, preverify_ok, crl=None):
+    """
+    验证对端证书的有效性
+    :param conn: OpenSSL.SSL.Connection, SSL 连接对象
+    :param cert: OpenSSL.crypto.X509, 当前证书
+    :param errno: int, OpenSSL错误代码, 0:无错误 | 9:证书过期 | 18: 自签名证书
+    :param depth: int, 当前证书在证书链中的深度 (0=叶子节点), 1:中间CA证书 -1:根CA证书 2+:更高级别CA证书
+    :param preverify_ok: int, 验证结果 (1=通过， 0=失败)
+    :param crl: _CRLInternal, CRL证书对象
+    :return: bool, True表示接受证书, False表示拒绝
+    """
+
+    if not preverify_ok:
+        from OpenSSL import SSL
+        error_str = SSL._ffi.string(SSL._lib.X509_verify_cert_error_string(errno)).decode()
+        logger.error(f"证书验证失败 (depth={depth}, err={errno}): {error_str}")
+        return False
+
+    if crl and is_certificate_revoked(cert, crl):
+        return False
+
+    return preverify_ok
+
+
+def load_ssl_pem(key_file, cert_file, ca_file, crl_file):
+    """
+    Load SSL PEM files.
+
+    Args:
+        key_file (str): The path to the private key file.
+        cert_file (str): The path to the certificate file.
+        ca_file (str): The path to the CA certificate file.
+        crl_file (str): The path to the CRL file.
+
+    Returns:
+        tuple: (key, crt, ca_crt, crl)
+
+    Raises:
+        Exception: If the file paths are invalid or the file contents are incorrect, exceptions may be thrown.
+    """
+
+    try:
+        # your_private_key_password
+        passphrase = ""
+        if not passphrase:
+            import pwinput
+            passphrase = pwinput.pwinput("Enter your password: ")
+        with FileOpen(key_file, "rb") as f:
+            key = crypto.load_privatekey(crypto.FILETYPE_PEM, f.read(), passphrase.encode())
+            del passphrase
+            gc.collect()
+        with FileOpen(cert_file, "rb") as f:
+            crt = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
+            check_crt_valid(crt)
+
+            crt_serial_number = hex(crt.get_serial_number())[2:]
+            logger.info(f"crt_serial_number: {crt_serial_number}")
+
+        check_certificate_match(crt, key)
+
+        with FileOpen(ca_file, "rb") as f:
+            ca_crt = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
+            check_crt_valid(ca_crt)
+
+            ca_serial_number = hex(ca_crt.get_serial_number())[2:]
+            logger.info(f"ca_serial_number: {ca_serial_number}")
+        crl = None
+        if os.path.exists(crl_file):
+            with FileOpen(crl_file, "rb") as f:
+                crl = x509.load_pem_x509_crl(f.read(), default_backend())
+                check_crl_valid(crl, ca_crt)
+            for revoked_cert in crl:
+                logger.info(f"Serial Number: {revoked_cert.serial_number}, "
+                            f"Revocation Date: {revoked_cert.revocation_date_utc}")
+
+    except Exception as e:
+        raise RuntimeError(f"The SSL certificate is invalid") from e
+
+    return key, crt, ca_crt, crl
+
+
+def check_crt_valid(pem):
+    """
+    Check the validity of the SSL certificate.
+
+    Raises:
+    RuntimeError: If the SSL certificate is invalid or expired.
+    """
+    try:
+        pem_start = parser.parse(pem.get_notBefore().decode("UTF-8"))
+        pem_end = parser.parse(pem.get_notAfter().decode("UTF-8"))
+        logger.info(f"The SSL certificate passes the verification and the validity period "
+                    f"starts from {pem_start} ends at {pem_end}.")
+    except Exception as e:
+        raise RuntimeError(f"The SSL certificate is invalid") from e
+
+    now_utc = datetime.now(tz=timezone.utc)
+    if pem.has_expired() or not (pem_start <= now_utc <= pem_end):
+        raise RuntimeError(f"The SSL certificate has expired.")
+
+
+def check_certificate_match(certificate, private_key):
+    """
+    Check certificate and private_key is match or not. if mismatched, an exception is thrown.
+    :param certificate:
+    :param private_key:
+    :return:
+    """
+    test_data = os.urandom(256)
+    try:
+        signature = crypto.sign(private_key, test_data, "sha256")
+        crypto.verify(
+            certificate,  # 包含公钥的证书
+            signature,  # 生成的签名
+            test_data,  # 原始数据
+            "sha256",  # 哈希算法
+        )
+        logger.info("公钥和私钥匹配")
+    except Exception as e:
+        raise RuntimeError("公钥和私钥不匹配") from e
+
+
+def check_crl_valid(crl, ca_crt):
+    # 验证CRL签名（确保CRL未被篡改）
+    if not crl.is_signature_valid(ca_crt.get_pubkey().to_cryptography_key()):
+        raise RuntimeError("CRL签名无效！")
+
+    # 检查CRL有效期
+    if not (crl.last_update <= datetime.utcnow() <= crl.next_update):
+        raise RuntimeError("CRL已过期或尚未生效!")

msprobe/pytorch/attl_manager.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025, 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.runtime import Runtime
+from msprobe.core.common.utils import Const
+from msprobe.pytorch.api_accuracy_checker.common.utils import ApiData
+from msprobe.pytorch.common.log import logger
+
+
+class ATTLManager:
+    def __init__(self, config):
+        self.config = config
+        self.attl = None
+
+    def attl_init(self):
+        if self.config.online_run_ut:
+            from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import ATTLConfig, ATTL
+            attl_config = ATTLConfig(is_benchmark_device=False,
+                                     connect_ip=self.config.host,
+                                     connect_port=self.config.port,
+                                     nfs_path=self.config.nfs_path,
+                                     tls_path=self.config.tls_path)
+            need_dump = len(self.config.rank) == 0 or Runtime.current_rank in self.config.rank
+            self.attl = ATTL('npu', attl_config, need_dump=need_dump)
+            if self.config.nfs_path:
+                self.attl.upload("start")
+
+    def attl_send(self, name, args, kwargs, output):
+        api_data = ApiData(
+                    name[:-len(Const.FORWARD_NAME_SUFFIX)],
+                    args,
+                    kwargs,
+                    output,
+                    Runtime.current_iter,
+                    Runtime.current_rank
+                )
+        logger.info(f"tools is dumping api: {api_data.name}, rank: {Runtime.current_rank}")
+        api_type, _, _ = api_data.name.split(Const.SEP)
+        if api_type in [Const.DISTRIBUTED]:
+            logger.info(f"api {api_data.name} is not supported, skip")
+            return
+        if self.config.nfs_path:
+            self.attl.upload(api_data)
+        else:
+            self.attl.send(api_data)
+
+    def attl_stop(self):
+        if self.config.nfs_path:
+            self.attl.upload("end")
+        elif self.attl.socket_manager is not None:
+            logger.info(f"pid: {os.getpid()} finished, start sends STOP signal.")
+            self.attl.socket_manager.send_stop_signal()

msprobe/pytorch/bench_functions/npu_fusion_attention.py

@@ -117,6 +117,12 @@ def fusion_attention_forward(forward_params):
     pse = forward_params.pse
     scale = forward_params.scale
     keep_prob = forward_params.keep_prob
+
+    # 除零风险拦截：keep_prob 为 0 时会导致除零错误
+    if keep_prob == 0:
+        raise ValueError("fusion_attention_forward: keep_prob cannot be zero to avoid division by zero.")
+
+
     qk = calculate_qk(q, k, atten_mask, pse, scale)
     softmax_res, softmax_max, softmax_sum = softmax_forward(qk)
     if drop_mask is None or len(drop_mask.shape) == 0:
@@ -137,6 +143,11 @@ def fusion_attention_backward(backward_params):
     pse = backward_params.pse
     scale = backward_params.scale
     keep_prob = backward_params.keep_prob
+
+    # 除零风险拦截：keep_prob 为 0 时会导致除零错误
+    if keep_prob == 0:
+        raise ValueError("fusion_attention_backward: keep_prob cannot be zero to avoid division by zero.")
+
     dp = torch.matmul(dx, v.permute(0, 1, 3, 2))
     if drop_mask is None or len(drop_mask.shape) == 0:
         drop_res = softmax_res.permute(0, 1, 3, 2)
@@ -164,23 +175,35 @@ def parse_bsnd_args(query, key, head_num, input_layout):
         if input_layout == "BSH":
             b, s1, h1 = query.shape
             _, s2, h2 = key.shape
+            if n1 == 0:
+                raise ValueError("parse_bsnd_args: head_num (n1) cannot be zero to avoid division by zero.")
             d = h1 // n1
+            if d == 0:
+                raise ValueError("parse_bsnd_args: computed head dimension (d) is zero, division by zero risk.")
             n2 = h2 // d
         elif input_layout == "SBH":
             s1, b, h1 = query.shape
             s2, _, h2 = key.shape
+            if n1 == 0:
+                raise ValueError("parse_bsnd_args: head_num (n1) cannot be zero to avoid division by zero.")
             d = h1 // n1
+            if d == 0:
+                raise ValueError("parse_bsnd_args: computed head dimension (d) is zero, division by zero risk.")
             n2 = h2 // d
         elif input_layout == "BSND":
             b, s1, n1, d = query.shape
             _, s2, n2, _ = key.shape
             h1 = n1 * d
             h2 = n2 * d
+            if d == 0:
+                raise ValueError("parse_bsnd_args: head dimension (d) is zero, division by zero risk.")
         elif input_layout == "BNSD":
             b, n1, s1, d = query.shape
             _, n2, s2, _ = key.shape
             h1 = n1 * d
             h2 = n2 * d
+            if d == 0:
+                raise ValueError("parse_bsnd_args: head dimension (d) is zero, division by zero risk.")
     except Exception as e:
         raise ValueError(f"query.shape: {query.shape}, key.shape: {key.shape}, parse_bsnd_args error: {e}") from e
 
@@ -446,6 +469,8 @@ def npu_fusion_attention_forward_patch(*args, **kwargs):
     input_layout = get_input_layout(*args, **kwargs)
 
     b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], head_num, input_layout)
+    if d == 0:
+        raise ValueError("npu_fusion_attention_forward_patch: head dimension (d) is zero, division by zero risk.")
     if n1 == n2 and s1 == s2:
         logger.debug(f"running case : BNSD = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
     else:
@@ -478,6 +503,8 @@ def npu_fusion_attention_backward_patch(*args, **kwargs):
         raise ValueError(f"Unsupported npu_fusion_attention_grad args {args}.")
 
     b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], args[4], args[5])
+    if d == 0:
+        raise ValueError("npu_fusion_attention_backward_patch: head dimension (d) is zero, division by zero risk.")
     if n1 == n2 and s1 == s2:
         logger.info(f"running case : bnsd = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
     else:

msprobe/pytorch/common/utils.py

@@ -24,6 +24,7 @@ from functools import wraps
 import numpy as np
 import torch
 import torch.distributed as dist
+
 from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import (FileCheckConst, change_mode,
                                             check_file_or_directory_path, check_path_before_create, FileOpen)
@@ -38,7 +39,9 @@ except ImportError:
 else:
     is_gpu = False
 
+
 torch_without_guard_version = torch.__version__ >= '2.1'
+torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 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
@@ -313,14 +316,14 @@ def print_rank_0(message):
         logger.info(message)
 
 
-def load_pt(pt_path, to_cpu=False):
+def load_pt(pt_path, to_cpu=False, weights_only=True):
     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"), weights_only=True)
+            pt = torch.load(pt_path, map_location=torch.device("cpu"), weights_only=weights_only)
         else:
-            pt = torch.load(pt_path, weights_only=True)
+            pt = torch.load(pt_path, weights_only=weights_only)
     except Exception as e:
         raise RuntimeError(f"load pt file {pt_path} failed") from e
     return pt
@@ -395,7 +398,7 @@ def save_api_data(api_data):
         io_buff = io.BytesIO()
         torch.save(api_data, io_buff)
     except Exception as e:
-        raise RuntimeError(f"save api_data to io_buff failed") from e
+        raise RuntimeError("save api_data to io_buff failed") from e
     return io_buff
 
 
@@ -403,9 +406,9 @@ def load_api_data(api_data_bytes):
     """Load data from bytes stream"""
     try:
         buffer = io.BytesIO(api_data_bytes)
-        buffer = torch.load(buffer, map_location="cpu")
+        buffer = torch.load(buffer, map_location="cpu", weights_only=False)
     except Exception as e:
-        raise RuntimeError(f"load api_data from bytes failed") from e
+        raise RuntimeError("load api_data from bytes failed") from e
     return buffer
 
 
@@ -457,7 +460,7 @@ def is_recomputation():
 
 def check_save_param(variable, name, save_backward):
     # try catch this api to skip invalid call
-    valid_data_types = tuple([torch.Tensor, int, float, str])
+    valid_data_types = (torch.Tensor, int, float, str)
     if not is_save_variable_valid(variable, valid_data_types):
         valid_data_types_with_nested_types = valid_data_types + (dict, tuple, list)
         logger.warning("PrecisionDebugger.save variable type not valid, "
@@ -476,13 +479,8 @@ def check_save_param(variable, name, save_backward):
         raise ValueError
 
 
-def replace_last_occurrence(text, old, new):
-    if text is None:
-        return text
-    index = text.rfind(old)
-    if index != -1:
-        return text[:index] + text[index:].replace(old, new, 1)
-    return text
+def is_torch_nn_module(variable):
+    return isinstance(variable, torch.nn.Module) and not isinstance(variable, torch.jit.ScriptModule)
 
 
 def is_hifloat8_tensor(tensor):
@@ -495,3 +493,17 @@ def is_float8_tensor(tensor):
     if str(tensor.dtype) in [Const.FLOAT8_E5M2_TYPE, Const.FLOAT8_E4M3FN_TYPE]:
         return True
     return is_hifloat8_tensor(tensor)
+
+
+def register_forward_pre_hook(module, forward_pre_hook):
+    if torch_version_above_or_equal_2:
+        module.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
+    else:
+        module.register_forward_pre_hook(forward_pre_hook)
+
+
+def register_forward_hook(module, forward_hook):
+    if torch_version_above_or_equal_2:
+        module.register_forward_hook(forward_hook, with_kwargs=True)
+    else:
+        module.register_forward_hook(forward_hook)

msprobe/pytorch/compare/distributed_compare.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-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");
@@ -13,41 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
-from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.common.file_utils import create_directory
-from msprobe.core.common.utils import CompareException, check_compare_param, check_configuration_param, get_dump_mode, \
-    set_dump_path
-from msprobe.core.compare.acc_compare import ModeConfig
-from msprobe.core.compare.utils import check_and_return_dir_contents, extract_json, set_stack_json_path
-from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.compare.pt_compare import PTComparator, compare
+from msprobe.core.compare.utils import compare_distributed_inner
+from msprobe.pytorch.compare.pt_compare import compare
 
 
 def compare_distributed(npu_dump_dir, bench_dump_dir, output_path, **kwargs):
-    if kwargs.get("suffix"):
-        logger.error("Argument 'suffix' is not supported for compare_distributed.")
-        raise CompareException(CompareException.INVALID_PARAM_ERROR)
-    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'))
-    if len(npu_ranks) != len(bench_ranks):
-        logger.error(
-            "The number of ranks in the two runs are different. "
-            "Unable to match the ranks. "
-            "Please use another folder to compare or use compare() api and manually match the ranks.")
-        raise CompareException(CompareException.INVALID_PATH_ERROR)
-    for nr, br in zip(npu_ranks, bench_ranks):
-        npu_data_dir = os.path.join(npu_dump_dir, nr)
-        bench_data_dir = os.path.join(bench_dump_dir, br)
-        npu_path = extract_json(npu_data_dir, stack_json=False)
-        bench_path = extract_json(bench_data_dir, stack_json=False)
-
-        dump_result_param = {
-            "npu_json_path": npu_path,
-            "bench_json_path": bench_path,
-            "is_print_compare_log": is_print_compare_log
-        }
-        compare(input_param=dump_result_param, output_path=output_path, suffix=f'_{nr}-{br}', **kwargs)
+    compare_distributed_inner(npu_dump_dir, bench_dump_dir, output_path, compare, **kwargs)