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

msprobe/core/config_check/ckpt_compare/megatron_loader.py

@@ -0,0 +1,302 @@
+# 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
+import re
+from collections import defaultdict
+from typing import Dict
+import numpy as np
+from msprobe.core.common.log import logger
+from msprobe.core.common.decorator import recursion_depth_decorator
+from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import FileOpen, load_yaml
+from msprobe.core.common.framework_adapter import FmkAdp
+
+# both weights and bias are partitioned in column parallel
+COLUMN_PARALLEL_PARAMS = ['linear_qkv', 'linear_fc1', 'word_embeddings.weight', 'output_layer.weight'] 
+# only weights are partitioned in column parallel
+ROW_PARALLEL_PARAMS = ['linear_fc2.weight', 'linear_proj.weight']
+ARGS = 'args'
+LAYER_IDX_PATTERN = re.compile('layers\.(\d+)\.')
+EXPERT_IDX_PATTERN = re.compile('experts\.(\d+)\.')
+ITER_DIR_PATTERN = re.compile('iter_([\d]{7})')
+
+
+@recursion_depth_decorator('')
+def _get_parameter(weights, prefix=''):
+    for k, v in weights.items():
+        name = Const.SEP.join([prefix, k]).strip(Const.SEP)
+        if isinstance(v, dict):
+            yield from _get_parameter(v, prefix=name)
+        elif FmkAdp.is_tensor(v):
+            yield name, FmkAdp.asnumpy(v)
+
+
+def _map_to_mcore_local_names(param_name: str) -> str:
+    """Map parameter names to mcore + local transformer implementation names."""
+    mcore_local_map = load_yaml(os.path.join(os.path.dirname(__file__), 'name_mapping.yaml'))
+    for other_name, mcore_local_name in mcore_local_map.items():
+        param_name = param_name.replace(other_name, mcore_local_name)
+    
+    return param_name
+
+
+def _parse_real_layer_idx(param_name, num_layers_per_stage, pp_size, pp_rank):
+    """Map local (virtual) pipeline stage layer index to global layer index.
+    
+    For virtual pipeline parallel, each pipeline stage is further divided into virtual stages.
+    The global layer index needs to account for both pipeline stage and virtual stage.
+    
+    Args:
+        param_name (str): Parameter name containing layer index: layers.x.<submodule_name>/<vpp_stage>
+        num_layers_per_stage (int): Number of layers per pipeline stage
+        pp_size (int): Pipeline parallel size
+        
+    Returns:
+        int: Global layer index accounting for both pipeline and virtual pipeline stages
+    """
+    # Extract local layer index from parameter name
+    layer_match = re.search(LAYER_IDX_PATTERN, param_name)
+    param_name, vpp_stage = param_name.split(Const.SCOPE_SEPARATOR)
+    if not layer_match:
+        return param_name
+    
+    local_layer_idx = int(layer_match.group(1))
+    vpp_stage = int(vpp_stage)
+    
+    # Calculate global layer index based on pipeline stage and virtual stage
+    real_layer_idx = local_layer_idx + (pp_size * vpp_stage + pp_rank) * num_layers_per_stage
+    
+    return param_name.replace(f'layers.{local_layer_idx}', f'layers.{real_layer_idx}')
+
+
+def _parse_real_expert_idx(param_name, num_experts_per_rank, exp_rank):
+    """Map local expert index to global expert index. TODO: shared expert
+    
+    For expert parallel, experts are distributed across ranks. This function maps
+    the local expert index on a rank to its global index across all ranks.
+    
+    Args:
+        param_name (str): Parameter name containing local expert index
+        num_experts_per_rank (int): Number of experts on each rank
+        exp_rank (int): Expert parallel rank
+        
+    Returns:
+        str: Parameter name with local expert index replaced by global expert index
+    """
+    # Extract local layer index from parameter name
+    expert_match = re.search(EXPERT_IDX_PATTERN, param_name)
+    if not expert_match:
+        return param_name
+    
+    local_expert_idx = int(expert_match.group(1))
+    # Calculate global layer index based on pipeline stage and virtual stage
+    real_experts_idx = local_expert_idx + exp_rank * num_experts_per_rank
+    
+    return param_name.replace(f'experts.{local_expert_idx}', f'experts.{real_experts_idx}')
+
+
+def _consolidate_tp_weights(weights: Dict) -> Dict:
+    """Consolidate weights from different tensor parallel ranks into combined tensors.
+    
+    Args:
+        weights: Dictionary of weights with rank information in keys
+        
+    Returns:
+        Dict: Consolidated weights without rank information
+    """
+    consolidated = {}
+    for key, tensors in weights.items():    
+        if any([name in key for name in COLUMN_PARALLEL_PARAMS]):
+            # Column parallel - concatenate along input dimension (dim 0)
+            combined = np.concatenate(tensors, axis=0)
+        elif any([name in key for name in ROW_PARALLEL_PARAMS]):
+            # Row parallel - concatenate along output dimension (dim 1)
+            combined = np.concatenate(tensors, axis=1)
+        else:
+            # For other params, verify identical and use first
+            if not all(np.allclose(tensors[0], t) for t in tensors[1:]):
+                logger.warning(f"Inconsistent values for {key} across TP ranks")
+            combined = tensors[0]
+    
+        consolidated[key] = combined
+    return consolidated
+
+
+def _parse_num_layers_per_stage(tp_partition):
+    match = [re.findall(LAYER_IDX_PATTERN, key) for key in tp_partition.keys()]
+    layer_idx = [int(i[0]) for i in match if i]
+    num_layers_per_pipeline_stage = max(layer_idx) + 1
+
+    return num_layers_per_pipeline_stage
+
+
+def parse_parallel_size(checkpoint_dir: str):
+    """Parse tensor, pipeline and expert parallel sizes from checkpoint filenames.
+    
+    Args:
+        checkpoint_dir (str): Directory containing checkpoint files
+        
+    Returns:
+        Namespace
+    """
+    # Find all rank directories
+    rank_dirs = [d for d in os.listdir(checkpoint_dir) if d.startswith('mp_rank_')]
+    
+    if not rank_dirs:
+        raise ValueError(f"No checkpoint rank directories found in {checkpoint_dir}")
+    
+    ckpt = FmkAdp.load_checkpoint(
+        os.path.join(checkpoint_dir, rank_dirs[0], 'model_optim_rng.pt'),
+        to_cpu=True,
+        weights_only=False)
+    args = ckpt[ARGS]
+    return (
+        args.tensor_model_parallel_size, 
+        args.pipeline_model_parallel_size, 
+        args.expert_model_parallel_size, 
+        args.num_experts
+    )
+
+
+def parse_iteration(checkpoint_path: str) -> Dict:
+    """
+    Parse the checkpoint iteration directory from a given checkpoint path.
+
+    If the path is a top-level checkpoint directory, this function reads the
+    'latest_checkpointed_iteration.txt' file to determine the latest iteration.
+    If the path is already an iteration directory (e.g., 'iter_0000005'), it extracts
+    the iteration number from the path.
+
+    Args:
+        checkpoint_path (str): Path to the checkpoint directory or iteration directory.
+
+    Returns:
+        str: The full path to the checkpoint directory for the determined iteration.
+
+    Raises:
+        ValueError: If the checkpoint directory for the determined iteration does not exist.
+    """
+    iteration = None
+    tracker_file = os.path.join(checkpoint_path, "latest_checkpointed_iteration.txt")
+    if os.path.exists(tracker_file):
+        with FileOpen(tracker_file, 'r') as f:
+            latest_iteration = f.read().strip()
+            if latest_iteration != 'release':
+                try:
+                    iteration = int(latest_iteration)
+                except Exception:
+                    logger.warning(
+                        f"The latest_checkpointed_iteration is supposed to be `release` or an int. \
+                        But {latest_iteration} is found."
+                    )
+            checkpoint_path = os.path.join(checkpoint_path, f'iter_{iteration:07d}')
+    else:
+        match = re.findall(ITER_DIR_PATTERN, checkpoint_path)
+        if match:
+            iteration = int(match[0])
+
+    # Checkpoint directory for this iteration
+    logger.info(f"Loaded checkpoint from iteration {iteration}")
+    
+    if not os.path.exists(checkpoint_path):
+        raise ValueError(f"Checkpoint directory not found: {checkpoint_path}")
+    
+    return checkpoint_path
+
+
+def get_weights_from_state_dict(state_dict):
+    weights = {}
+    vpp_stage = 0
+    if 'model' in state_dict:
+        model_weights = state_dict['model']
+
+        for key, value in _get_parameter(model_weights):
+            key = _map_to_mcore_local_names(key)
+            weights[f"{key}{Const.SCOPE_SEPARATOR}{vpp_stage}"] = value
+
+    elif 'model0' in state_dict:
+        #vpp enabled
+        while f'model{vpp_stage}' in state_dict:
+            model_weights = state_dict[f'model{vpp_stage}']
+            for key, value in _get_parameter(model_weights):
+                key = _map_to_mcore_local_names(key)
+                weights[f"{key}{Const.SCOPE_SEPARATOR}{vpp_stage}"] = value
+            vpp_stage += 1
+    return weights
+
+
+def load_megatron_weights(checkpoint_path: str) -> Dict:
+    """Load Megatron parallel checkpoint weights into a single dictionary.
+    
+    Args:
+        checkpoint_path (str): Base checkpoint directory path
+
+    Returns:
+        combined_weights: Dict with weights from all ranks, keys include rank info
+    """
+    try:
+        import megatron
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError("No module named 'megatron', which is required to load a megatron ckpt") from e
+
+    # Find latest iteration if not specified
+    checkpoint_path = parse_iteration(checkpoint_path)
+
+    # Parse parallel sizes from checkpoint directory structure
+    tp_size, pp_size, exp_size, num_experts = parse_parallel_size(checkpoint_path)
+    combined_weights = {}
+
+    # Load checkpoints from all ranks
+    for exp_rank in range(exp_size):
+        num_layers_per_pipeline_stage = 0
+        for pp_rank in range(pp_size):
+            tp_partition = defaultdict(list)
+            for tp_rank in range(tp_size):
+                # Construct checkpoint path based on parallel ranks
+                if pp_size > 1:
+                    rank_dir = f'mp_rank_{tp_rank:02d}_{pp_rank:03d}'
+                else:
+                    rank_dir = f'mp_rank_{tp_rank:02d}'
+                
+                if exp_size > 1:
+                    rank_dir = f'{rank_dir}_{exp_rank:03d}'
+                
+                ckpt_file = os.path.join(checkpoint_path, rank_dir, 'model_optim_rng.pt')
+                try:
+                    state_dict = FmkAdp.load_checkpoint(ckpt_file, to_cpu=True, weights_only=False)
+                    partition = get_weights_from_state_dict(state_dict)
+                    for key, weight in partition.items():
+                        tp_partition[key].append(weight)
+                    
+                except Exception as load_error:
+                    logger.warning(f"Error loading {ckpt_file}: {load_error}")
+            
+            if not tp_partition:
+                raise ValueError('No state loaded.')
+            
+            if not num_layers_per_pipeline_stage:
+                num_layers_per_pipeline_stage = _parse_num_layers_per_stage(tp_partition)
+
+            consolidated_weight = _consolidate_tp_weights(tp_partition)
+            for key, value in consolidated_weight.items(): 
+                key = _parse_real_layer_idx(key, num_layers_per_pipeline_stage, pp_size, pp_rank)
+                if num_experts:
+                    key = _parse_real_expert_idx(key, num_experts // exp_size, exp_rank)
+                combined_weights[key] = value
+
+    logger.info(f"Found {len(combined_weights)} total parameters across all ranks")
+
+    return combined_weights

msprobe/core/config_check/ckpt_compare/metrics.py

@@ -0,0 +1,83 @@
+# 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 numpy as np
+
+from msprobe.core.common.log import logger
+from msprobe.core.compare.npy_compare import CompareOps
+
+
+
+def in_different_shape(a, b):
+    if a.shape != b.shape:
+        logger.warning(f"a, b are in different shape. a: {a.shape}, b: {b.shape}")
+        return True
+    return False
+
+
+def l2_distance(a, b):
+    if a is None or b is None:
+        return None
+    if in_different_shape(a, b):
+        return None
+    return np.linalg.norm(a - b).item()
+
+
+def cos_sim(a, b):
+    if a is None or b is None:
+        return None
+
+    if in_different_shape(a, b):
+        return None
+    if a.ndim > 0:
+        a = a.flatten().squeeze()
+        b = b.flatten().squeeze()
+
+    num = a.dot(b)
+    a_norm = np.linalg.norm(a)
+    b_norm = np.linalg.norm(b)
+
+    if a_norm == 0 and b_norm == 0:
+        return 1.
+    if a_norm == 0 or b_norm == 0:
+        logger.warning(f'One tensor norm is zero.')
+        return None
+
+    sim = num / (a_norm * b_norm)
+
+    return sim.item()
+
+
+def numel(a, b):
+    n1 = a.size
+    n2 = b.size
+    if n1 != n2:
+        logger.warning('parameters have different number of element')
+        return (n1, n2)
+    return n1
+
+
+def shape(a, b):
+    if in_different_shape(a, b):
+        return [list(a.shape), list(b.shape)]
+    return list(a.shape)
+
+
+METRIC_FUNC = {
+    'l2': l2_distance, 
+    'cos': cos_sim, 
+    'numel': numel, 
+    'shape': shape
+    }

msprobe/core/config_check/ckpt_compare/name_mapping.yaml

@@ -0,0 +1,12 @@
+self_attention.linear_qkv.layer_norm_: input_layernorm.
+language_model.: ''
+encoder: decoder
+.input_norm.: .input_layernorm.
+query_key_value: linear_qkv
+.dense.: .linear_proj.
+post_attention_norm: pre_mlp_layernorm
+dense_h_to_4h: linear_fc1
+dense_4h_to_h: linear_fc2
+mlp.local_experts: mlp.experts.local_experts
+final_norm: final_layernorm
+word_embeddings_for_head: output_layer

msprobe/core/config_check/config_check_cli.py

@@ -0,0 +1,51 @@
+# 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.config_check.config_checker import ConfigChecker
+from msprobe.core.config_check.ckpt_compare.ckpt_comparator import compare_checkpoints
+from msprobe.core.common.log import logger
+
+
+def pack(shell_path, output_path, framework):
+    ConfigChecker(shell_path=shell_path, output_zip_path=output_path, fmk=framework)
+
+
+def compare(bench_zip_path, cmp_zip_path, output_path, framework):
+    ConfigChecker.compare(bench_zip_path, cmp_zip_path, output_path, framework)
+
+
+def _config_checking_parser(parser):
+    parser.add_argument('-d', '--dump', nargs='*', help='Collect the train config into a zip file')
+    parser.add_argument('-c', '--compare', nargs=2, help='Compare two zip files or checkpoints')
+    parser.add_argument('-o', '--output', help='output path, default is current directory')
+
+
+def _run_config_checking_command(args):
+    if args.dump is not None:
+        output_dirpath = args.output if args.output else "./config_check_pack.zip"
+        pack(args.dump, output_dirpath, args.framework)
+    elif args.compare:
+        if args.compare[0].endswith('zip'):
+            logger.info('The input paths is zip files, comparing packed config.')
+            output_dirpath = args.output if args.output else "./config_check_result"
+            compare(args.compare[0], args.compare[1], output_dirpath, args.framework)
+        else:
+            logger.info('Comparing model checkpoint.')
+            output_dirpath = args.output if args.output else "./ckpt_similarity.json"
+            compare_checkpoints(args.compare[0], args.compare[1], output_dirpath)
+
+    else:
+        logger.error("The param is not correct, you need to give '-d' for dump or '-c' for compare.")
+        raise Exception("The param is not correct, you need to give '-d' for dump or '-c' for compare.")

msprobe/core/config_check/config_checker.py

@@ -0,0 +1,100 @@
+# 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
+import shutil
+
+import pandas as pd
+
+from msprobe.core.common.file_utils import save_excel, split_zip_file_path, \
+    create_directory, extract_zip
+from msprobe.core.common.framework_adapter import FmkAdp
+from msprobe.core.config_check.checkers.base_checker import PackInput
+from msprobe.core.config_check.utils.utils import config_checking_print
+from msprobe.core.common.const import Const
+
+
+class ConfigChecker:
+    checkers = {}
+    pre_forward_fun_list = []
+    result_filename = "result.xlsx"
+    result_header = ["filename", "pass_check"]
+    step = 0
+
+    def __init__(self, model=None, shell_path=None, output_zip_path="./config_check_pack.zip", fmk="pytorch"):
+        FmkAdp.set_fmk(fmk)
+        self.pack_input = PackInput(output_zip_path, model, shell_path)
+        file_path, file_name = split_zip_file_path(self.pack_input.output_zip_path)
+        if not os.path.exists(file_path):
+            create_directory(file_path)
+        self.pack()
+
+    @staticmethod
+    def compare(bench_zip_path, cmp_zip_path, output_path, fmk=Const.PT_FRAMEWORK):
+        if os.path.exists(output_path):
+            shutil.rmtree(output_path)
+        bench_dir = os.path.join(output_path, "bench")
+        cmp_dir = os.path.join(output_path, "cmp")
+        extract_zip(bench_zip_path, bench_dir)
+        config_checking_print(f"extract zip file {bench_zip_path} to {bench_dir}")
+        extract_zip(cmp_zip_path, cmp_dir)
+        config_checking_print(f"extract zip file {cmp_zip_path} to {cmp_dir}")
+
+        result = []
+        summary_result = []
+        for checker in ConfigChecker.checkers.values():
+            checker_name, pass_check, df = checker.compare_ex(bench_dir, cmp_dir, output_path, fmk)
+            if checker_name:
+                summary_result.append([checker_name, pass_check])
+            if df is not None:
+                result.append((df, checker_name))
+        summary_result_df = pd.DataFrame(summary_result, columns=ConfigChecker.result_header)
+        result.insert(0, (summary_result_df, "summary"))
+        save_excel(os.path.join(output_path, ConfigChecker.result_filename), result)
+        config_checking_print(f"config checking result save to {os.path.realpath(output_path)}")
+
+    @staticmethod
+    def apply_patches(fmk=Const.PT_FRAMEWORK):
+        for checker in ConfigChecker.checkers.values():
+            checker.apply_patches(fmk)
+
+    def pack(self):
+        config_checking_print(f"pack result zip path {os.path.realpath(self.pack_input.output_zip_path)}")
+
+        def hook(model, args, kwargs):
+            for collect_func in self.pre_forward_fun_list:
+                collect_func(model, args, kwargs, ConfigChecker.step)
+            ConfigChecker.step += 1
+
+        if self.pack_input.model:
+            FmkAdp.register_forward_pre_hook(self.pack_input.model, hook, with_kwargs=True)
+        for checker in ConfigChecker.checkers.values():
+            if checker.input_needed and not getattr(self.pack_input, checker.input_needed):
+                continue
+            if FmkAdp.is_initialized() and FmkAdp.get_rank() != 0 and not checker.multi_rank:
+                continue
+            checker.pack(self.pack_input)
+
+
+def register_checker_item(key, cls=None):
+    if cls is None:
+        # 无参数时，返回装饰器函数
+        return lambda cls: register_checker_item(key, cls)
+    ConfigChecker.checkers[key] = cls
+    return cls
+
+
+def register_pre_forward_fun_list(func):
+    ConfigChecker.pre_forward_fun_list.append(func)

msprobe/{pytorch/parse.py → core/config_check/resource/dependency.yaml}

@@ -13,7 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from msprobe.pytorch.parse_tool import cli
-
-if __name__ == '__main__':
-    cli.parse()
+dependency:
+  - transformers
+  - deepspeed
+  - megatron
+  - numpy
+  - datasets
+  - peft

msprobe/core/config_check/resource/env.yaml

@@ -0,0 +1,57 @@
+# 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.
+
+HCCL_DETERMINISTIC:
+  npu:
+    name: HCCL_DETERMINISTIC
+    default_value: False
+  gpu:
+    name: NCCL_DETERMINISTIC
+    default_value: False
+
+HCCL_ALGO:
+  npu:
+    name: HCCL_ALGO
+    default_value: None
+  gpu:
+    name: NCCL_ALGO
+    default_value: None
+
+HCCL_INTRA_ROCE_ENABLE:
+  npu:
+    name: HCCL_INTRA_ROCE_ENABLE
+    default_value: 0
+
+
+HCCL_INTRA_PICE_ENABLE:
+  npu:
+    name: HCCL_INTRA_ROCE_ENABLE
+    default_value: 1
+
+ASCEND_LAUNCH_BLOCKING:
+  npu:
+    name: ASCEND_LAUNCH_BLOCKING
+    default_value: 0
+  gpu:
+    name: CUDA_LAUNCH_BLOCKING
+    default_value: 0
+
+ASCEND_RT_VISIBLE_DEVICES:
+  npu:
+    name: ASCEND_RT_VISIBLE_DEVICES
+    default_value: None
+  gpu:
+    name: CUDA_VISIBLE_DEVICES
+    default_value: None

msprobe/core/config_check/resource/hyperparameter.yaml

@@ -0,0 +1,21 @@
+learning_rate:
+  - lr
+  - learningrate
+
+batch_size:
+  - batch
+  - bs
+  - batch_size_per_gpu
+
+epochs:
+  - num_epochs
+  - max_epochs
+  - epoch
+
+weight_decay:
+  - wd
+  - weightdecay
+
+dropout_rate:
+  - dropout
+  - drop_rate