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

msprobe/pytorch/bench_functions/rms_norm.py

@@ -1,15 +1,15 @@
-import torch
-
-
-def npu_rms_norm(x, gamma, epsilon=1e-5):
-    rstd = torch.rsqrt(torch.mean(torch.pow(x, 2), axis=-1, keepdim=True) + epsilon)
-    res = x * rstd * gamma
-    return res.cpu(), rstd.float().cpu()
-
-
-def npu_rms_norm_backward(grad, x, gamma, rstd):
-    mean_gy = (grad * x * gamma * rstd).mean(dim=-1, keepdim=True)
-    grad_x = (grad * gamma - x * rstd * mean_gy) * rstd
-    grad_gamma = x * grad * rstd
-    return grad_x.cpu(), grad_gamma.cpu()
-
+import torch
+
+
+def npu_rms_norm(x, gamma, epsilon=1e-5):
+    rstd = torch.rsqrt(torch.mean(torch.pow(x, 2), axis=-1, keepdim=True) + epsilon)
+    res = x * rstd * gamma
+    return res, rstd.float()
+
+
+def npu_rms_norm_backward(grad, x, gamma, rstd):
+    mean_gy = (grad * x * gamma * rstd).mean(dim=-1, keepdim=True)
+    grad_x = (grad * gamma - x * rstd * mean_gy) * rstd
+    grad_gamma = x * grad * rstd
+    return grad_x.cpu(), grad_gamma.cpu()
+

msprobe/pytorch/bench_functions/rotary_mul.py

@@ -1,52 +1,52 @@
-import torch
-
-
-def npu_rotary_mul(x, r1, r2):
-    x1, x2 = torch.chunk(x, 2, -1)
-    x_new = torch.cat((-x2, x1), dim=-1)
-    output = r1 * x + r2 * x_new
-    return output.cpu()
-
-
-def npu_rotary_mul_backward(dy_tensor, x, r1, r2):
-    x.requires_grad = True
-    r1.requires_grad = True
-    r2.requires_grad = True
-    # golden
-    x1, x2 = torch.chunk(x, 2, -1)
-    x_new = torch.cat((-x2, x1), dim=-1)
-    golden_tensor = r1 * x + r2 * x_new
-    golden_tensor.backward(dy_tensor)
-    r1_shape = r1.shape
-    r1_grad = torch.zeros(r1_shape).type(torch.float32)
-    r2_grad = torch.zeros(r1_shape).type(torch.float32)
-    x1, x2 = torch.chunk(x.float(), 2, -1)
-    x_new2 = torch.cat((-x2, x1), dim=-1)
-    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 condition_1:
-        for i in range(x_shape[0]):
-            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, :])
-    elif condition_2:
-        for i in range(x_shape[0]):
-            for j in range(x_shape[1]):
-                r2_grad[0, 0, :, :] += (x_new2[i, j, :, :] * grad[i, j, :, :])
-                r1_grad[0, 0, :, :] += (h[i, j, :, :] * grad[i, j, :, :])
-    elif condition_3:
-        for i in range(x_shape[1]):
-            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()
+import torch
+
+
+def npu_rotary_mul(x, r1, r2):
+    x1, x2 = torch.chunk(x, 2, -1)
+    x_new = torch.cat((-x2, x1), dim=-1)
+    output = r1 * x + r2 * x_new
+    return output
+
+
+def npu_rotary_mul_backward(dy_tensor, x, r1, r2):
+    x.requires_grad = True
+    r1.requires_grad = True
+    r2.requires_grad = True
+    # golden
+    x1, x2 = torch.chunk(x, 2, -1)
+    x_new = torch.cat((-x2, x1), dim=-1)
+    golden_tensor = r1 * x + r2 * x_new
+    golden_tensor.backward(dy_tensor)
+    r1_shape = r1.shape
+    r1_grad = torch.zeros(r1_shape).type(torch.float32)
+    r2_grad = torch.zeros(r1_shape).type(torch.float32)
+    x1, x2 = torch.chunk(x.float(), 2, -1)
+    x_new2 = torch.cat((-x2, x1), dim=-1)
+    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 condition_1:
+        for i in range(x_shape[0]):
+            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, :])
+    elif condition_2:
+        for i in range(x_shape[0]):
+            for j in range(x_shape[1]):
+                r2_grad[0, 0, :, :] += (x_new2[i, j, :, :] * grad[i, j, :, :])
+                r1_grad[0, 0, :, :] += (h[i, j, :, :] * grad[i, j, :, :])
+    elif condition_3:
+        for i in range(x_shape[1]):
+            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,26 +1,26 @@
-import torch
-
-
-def npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask):
-    if fixed_triu_mask:
-        mask = (torch.triu(torch.ones(mask.shape), k=1)).bool().to(mask.device)
-    dtype = x.dtype
-    x = (x * scale).masked_fill(mask, value=-10000)
-    x = x - torch.max(x, dim=-1, keepdims=True)[0]
-    x = torch.exp(x.float())
-    y = torch.div(x, torch.sum(x, dim=-1, keepdims=True))
-    return y.to(dtype).cpu()
-
-
-def npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask):
-    if fixed_triu_mask:
-        mask = (torch.triu(torch.ones(mask.shape), k=1)).bool().to(mask.device)
-    dtype = y_grad.dtype
-    y_grad = y_grad.float()
-    y = y.float()
-    x_grad = y_grad * y
-    x_grad = y_grad - torch.sum(x_grad, dim=-1, keepdims=True)
-    x_grad = x_grad * y
-    x_grad = x_grad * scale
-    x_grad = x_grad.masked_fill(mask, value=0)
-    return x_grad.to(dtype).cpu()
+import torch
+
+
+def npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask):
+    if fixed_triu_mask:
+        mask = (torch.triu(torch.ones(mask.shape), k=1)).bool().to(mask.device)
+    dtype = x.dtype
+    x = (x * scale).masked_fill(mask, value=-10000)
+    x = x - torch.max(x, dim=-1, keepdims=True)[0]
+    x = torch.exp(x.float())
+    y = torch.div(x, torch.sum(x, dim=-1, keepdims=True))
+    return y.to(dtype)
+
+
+def npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask):
+    if fixed_triu_mask:
+        mask = (torch.triu(torch.ones(mask.shape), k=1)).bool().to(mask.device)
+    dtype = y_grad.dtype
+    y_grad = y_grad.float()
+    y = y.float()
+    x_grad = y_grad * y
+    x_grad = y_grad - torch.sum(x_grad, dim=-1, keepdims=True)
+    x_grad = x_grad * y
+    x_grad = x_grad * scale
+    x_grad = x_grad.masked_fill(mask, value=0)
+    return x_grad.to(dtype).cpu()

msprobe/pytorch/bench_functions/swiglu.py

@@ -1,55 +1,55 @@
-import torch
-
-
-def npu_swiglu(x, dim=-1):
-    tensor_dtype = x.dtype
-
-    inTensors = torch.chunk(x, 2, dim=dim)
-    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])
-    else:
-        tensor_self_float = inTensors[0].type(torch.float)
-        tensor_other_float = inTensors[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)
-    return output_data.cpu()
-
-
-def npu_swiglu_backward(grad, x, dim=-1):
-    tensor_dtype = grad.dtype
-    in_tensors = torch.chunk(x, 2, dim=dim)
-    tensor_grad_out = grad
-
-    if tensor_dtype == torch.float16:
-        tensor_out1 = torch.mul(
-            torch.mul(in_tensors[1].type(torch.float32), swish_grad(1, in_tensors[0].type(torch.float32))),
-            tensor_grad_out.type(torch.float32)).type(torch.float16)
-        tensor_out2 = torch.mul(tensor_grad_out.type(torch.float32),
-                                swish(1, in_tensors[0].type(torch.float32))).type(torch.float16)
-        output = torch.cat((tensor_out1, tensor_out2), dim)
-    elif tensor_dtype == torch.bfloat16:
-        tensor_self_float = in_tensors[0].type(torch.float)
-        tensor_other_float = in_tensors[1].type(torch.float)
-        tensor_gradout_float = tensor_grad_out.type(torch.float)
-
-        tensor_out1 = torch.mul(tensor_gradout_float, swish_grad(1.0, tensor_self_float)).type(torch.bfloat16).type(
-            torch.float32) * tensor_other_float
-        tensor_out2 = swish(1.0, tensor_self_float).type(torch.bfloat16).type(torch.float32) * tensor_gradout_float
-        tensor_out_float = torch.cat((tensor_out1, tensor_out2), dim=dim)
-        output = tensor_out_float.type(torch.bfloat16)
-    else:
-        tensor_out1 = torch.mul(torch.mul(in_tensors[1], swish_grad(1.0, in_tensors[0])), tensor_grad_out)
-        tensor_out2 = torch.mul(tensor_grad_out, swish(1.0, in_tensors[0]))
-        output = torch.cat((tensor_out1, tensor_out2), dim)
-    return output.cpu()
-
-
-def swish_grad(beta, x):
-    return torch.sigmoid(beta * x) + x * (1 - torch.sigmoid(beta * x)) * torch.sigmoid(beta * x) * beta
-
-
-def swish(beta, x):
-    return x * torch.sigmoid(beta * x)
-
+import torch
+
+
+def npu_swiglu(x, dim=-1):
+    tensor_dtype = x.dtype
+
+    inTensors = torch.chunk(x, 2, dim=dim)
+    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])
+    else:
+        tensor_self_float = inTensors[0].type(torch.float)
+        tensor_other_float = inTensors[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)
+    return output_data
+
+
+def npu_swiglu_backward(grad, x, dim=-1):
+    tensor_dtype = grad.dtype
+    in_tensors = torch.chunk(x, 2, dim=dim)
+    tensor_grad_out = grad
+
+    if tensor_dtype == torch.float16:
+        tensor_out1 = torch.mul(
+            torch.mul(in_tensors[1].type(torch.float32), swish_grad(1, in_tensors[0].type(torch.float32))),
+            tensor_grad_out.type(torch.float32)).type(torch.float16)
+        tensor_out2 = torch.mul(tensor_grad_out.type(torch.float32),
+                                swish(1, in_tensors[0].type(torch.float32))).type(torch.float16)
+        output = torch.cat((tensor_out1, tensor_out2), dim)
+    elif tensor_dtype == torch.bfloat16:
+        tensor_self_float = in_tensors[0].type(torch.float)
+        tensor_other_float = in_tensors[1].type(torch.float)
+        tensor_gradout_float = tensor_grad_out.type(torch.float)
+
+        tensor_out1 = torch.mul(tensor_gradout_float, swish_grad(1.0, tensor_self_float)).type(torch.bfloat16).type(
+            torch.float32) * tensor_other_float
+        tensor_out2 = swish(1.0, tensor_self_float).type(torch.bfloat16).type(torch.float32) * tensor_gradout_float
+        tensor_out_float = torch.cat((tensor_out1, tensor_out2), dim=dim)
+        output = tensor_out_float.type(torch.bfloat16)
+    else:
+        tensor_out1 = torch.mul(torch.mul(in_tensors[1], swish_grad(1.0, in_tensors[0])), tensor_grad_out)
+        tensor_out2 = torch.mul(tensor_grad_out, swish(1.0, in_tensors[0]))
+        output = torch.cat((tensor_out1, tensor_out2), dim)
+    return output.cpu()
+
+
+def swish_grad(beta, x):
+    return torch.sigmoid(beta * x) + x * (1 - torch.sigmoid(beta * x)) * torch.sigmoid(beta * x) * beta
+
+
+def swish(beta, x):
+    return x * torch.sigmoid(beta * x)
+

msprobe/pytorch/common/__init__.py

@@ -1,2 +1,2 @@
-from .parse_json import parse_json_info_forward_backward
-from .utils import seed_all
+from .parse_json import parse_json_info_forward_backward
+from .utils import seed_all

msprobe/pytorch/common/compare_script.template

@@ -1,14 +1,14 @@
-from ptdbg_ascend import compare
-
-pkl_path = "%s"
-dump_data_dir = "%s"
-
-dump_path_param = {
-    "npu_pkl_path": ,
-    "bench_pkl_path": ,
-    "npu_dump_data_dir": ,
-    "bench_dump_data_dir": ,
-    "is_print_compare_log": True
-}
-
-compare(dump_path_param, output_path="", stack_mode=%s)
+from ptdbg_ascend import compare
+
+pkl_path = "%s"
+dump_data_dir = "%s"
+
+dump_path_param = {
+    "npu_pkl_path": ,
+    "bench_pkl_path": ,
+    "npu_dump_data_dir": ,
+    "bench_dump_data_dir": ,
+    "is_print_compare_log": True
+}
+
+compare(dump_path_param, output_path="", stack_mode=%s)

msprobe/pytorch/common/log.py

@@ -1,32 +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
-from msprobe.core.common.exceptions import DistributedNotInitializedError
-
-
-class PyTorchLogger(BaseLogger):
-    def __init__(self):
-        super().__init__()
-
-    def get_rank(self):
-        try:
-            current_rank = get_rank_if_initialized()
-        except DistributedNotInitializedError:
-            current_rank = None
-        return current_rank
-
-    def _print_log(self, level, msg, end='\n'):
-        current_rank = self.get_rank()
-        current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
-        pid = os.getpid()
-        if current_rank is not None:
-            full_msg = f"{current_time} ({pid}) [rank {current_rank}] [{level}] {msg}"
-        else:
-            full_msg = f"{current_time} ({pid}) [{level}] {msg}"
-        print(full_msg, end=end)
-        sys.stdout.flush()
-
-
+import os
+import time
+import sys
+from msprobe.pytorch.common.utils import get_rank_if_initialized
+from msprobe.core.common.log import BaseLogger
+from msprobe.core.common.exceptions import DistributedNotInitializedError
+
+
+class PyTorchLogger(BaseLogger):
+    def __init__(self):
+        super().__init__()
+
+    def get_rank(self):
+        try:
+            current_rank = get_rank_if_initialized()
+        except DistributedNotInitializedError:
+            current_rank = None
+        return current_rank
+
+
 logger = PyTorchLogger()

msprobe/pytorch/common/parse_json.py

@@ -1,39 +1,39 @@
-import json
-
-from msprobe.core.common.exceptions import ParseJsonException
-from msprobe.core.common.file_check import FileOpen
-
-
-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)
-
-    real_data_path = dump_json.get("dump_data_dir")
-    dump_data = dump_json.get("data")
-    if not dump_data:
-        raise ParseJsonException(ParseJsonException.InvalidDumpJson, "dump数据中没有data字段")
-
-    forward_data = {}
-    backward_data = {}
-    for data_name, data_item in dump_data.items():
-        if "Module" in data_name:
-            continue
-        if "forward" in data_name:
-            api_name = parse_data_name_with_pattern(data_name, "forward")
-            forward_data.update({api_name: data_item})
-        elif "backward" in data_name:
-            api_name = parse_data_name_with_pattern(data_name, "backward")
-            backward_data.update({api_name: data_item})
-        else:
-            raise ParseJsonException(ParseJsonException.UnexpectedNameStruct,
-                f"{data_name} in file {json_path}.")
-
-    return forward_data, backward_data, real_data_path
+import json
+
+from msprobe.core.common.exceptions import ParseJsonException
+from msprobe.core.common.file_utils import FileOpen
+
+
+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)
+
+    real_data_path = dump_json.get("dump_data_dir")
+    dump_data = dump_json.get("data")
+    if not dump_data:
+        raise ParseJsonException(ParseJsonException.InvalidDumpJson, "dump数据中没有data字段")
+
+    forward_data = {}
+    backward_data = {}
+    for data_name, data_item in dump_data.items():
+        if "Module" in data_name:
+            continue
+        if "forward" in data_name:
+            api_name = parse_data_name_with_pattern(data_name, "forward")
+            forward_data.update({api_name: data_item})
+        elif "backward" in data_name:
+            api_name = parse_data_name_with_pattern(data_name, "backward")
+            backward_data.update({api_name: data_item})
+        else:
+            raise ParseJsonException(ParseJsonException.UnexpectedNameStruct,
+                f"{data_name} in file {json_path}.")
+
+    return forward_data, backward_data, real_data_path