mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/scripts/run_script/main_msaccucmp.sh

@@ -0,0 +1,232 @@
+#!/bin/bash
+# install.sh - 安装脚本
+USERNAME=$(id -un)
+USERGROUP=$(id -gn)
+
+COMMON_SHELL_PATH="common.sh"
+
+# 颜色定义
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+# 打印彩色信息
+info() {
+    echo -e "${GREEN}[INFO]${NC} $1"
+}
+
+warn() {
+    echo -e "${YELLOW}[WARN]${NC} $1"
+}
+
+error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+# 检查是否已安装
+check_installation() {
+    if [[ -d "$INSTALL_DIR" ]]; then
+        if [[ "$SILENT_MODE" != "true" ]]; then
+            echo -n "目录 $INSTALL_DIR 已存在，是否覆盖？ (y/N): "
+            read -r response
+            if [[ ! "$response" =~ ^[Yy]$ ]]; then
+                info "安装已取消"
+                exit 0
+            fi
+        fi
+        warn "将覆盖现有安装: $INSTALL_DIR"
+    fi
+}
+
+# 创建安装目录
+create_install_dir() {
+    info "创建安装目录: $INSTALL_DIR"
+
+    # 尝试创建目录
+    if ! mkdir -p "$INSTALL_DIR"; then
+        error "目录创建失败: $INSTALL_DIR"
+        error "请检查是否有足够的权限，或手动创建该目录"
+        exit 1
+    fi
+}
+
+function convert_install_path() {
+    local _install_path="$1"
+
+    # delete last "/" "/."，并将中间的"/./"替换成"/"
+    _install_path=`echo "${_install_path}" | sed -r "s/((\/)|(\/\.))*$//g" | sed -r "s|/\./|/|g"`
+    if [ -z "${_install_path}" ]; then
+        _install_path="/"
+    fi
+    # covert relative path to absolute path
+    # 处理以 "./" 或 "." 开头的路径
+    if [[ "${_install_path}" =~ ^\./.* ]] || [[ "${_install_path}" =~ ^\.$ ]]; then
+        # 移除开头的 "./" 或 "."
+        _install_path=`echo "${_install_path}" | sed -r "s|^\./||" | sed -r "s|^\.$||"`
+        if [ -z "${_install_path}" ]; then
+            _install_path="${run_path}"
+        else
+            _install_path="${run_path}/${_install_path}"
+        fi
+    else
+        local _prefix=`echo "${_install_path}" | cut -d"/" -f1`
+        if [ ! -z "${_prefix}" ] && [ "~" != "${_prefix}" ]; then
+            _install_path="${run_path}/${_install_path}"
+        fi
+    fi
+    # covert '~' to home path
+    local _suffix_path=`echo "${_install_path}" | cut -d"~" -f2`
+    if [ "${_suffix_path}" != "${_install_path}" ]; then
+        local _home_path=`eval echo "~" | sed "s/\/*$//g"`
+        _install_path="${_home_path}${_suffix_path}"
+    fi
+    # 规范化路径：移除双斜杠，处理相对路径组件
+    _install_path=`echo "${_install_path}" | sed -r "s|//+|/|g"`
+    # 转换为绝对路径（如果可能）
+    if command -v realpath >/dev/null 2>&1; then
+        _install_path=$(realpath -m "${_install_path}" 2>/dev/null || echo "${_install_path}")
+    fi
+    echo "${_install_path}"
+}
+
+# 主安装函数
+main_install() {
+    # 设置安装目录
+    INSTALL_DIR=$(get_install_path)
+
+    if [[ "$SILENT_MODE" != "true" ]]; then
+        echo "========================================="
+        echo "  Ascend-mindstudio-accucmp 安装程序"
+        echo "========================================="
+        echo "安装目录: $INSTALL_DIR"
+        echo ""
+    fi
+
+    # 检查是否已安装
+    check_installation
+    # 创建安装目录
+    create_install_dir
+
+    info "正在安装文件..."
+
+    # 复制文件到安装目录
+    if [[ "$SUDO_USED" == "true" ]]; then
+        sudo cp -r "$PWD"/msaccucmp/* "$INSTALL_DIR"/
+    else
+        cp -r "$PWD"/msaccucmp/* "$INSTALL_DIR"/
+    fi
+
+    info "安装完成！"
+    info "安装目录: $INSTALL_DIR"
+}
+
+
+function get_install_path() {
+    if [ -z "${input_install_path}" ]; then
+        local _install_path
+        _install_path="/usr/local/Ascend/operator_cmp/compare"
+    else
+        _install_path="${input_install_path}/operator_cmp/compare"
+    fi
+    install_path=$(convert_install_path "${_install_path}")
+    echo "${install_path}"
+}
+
+# 主函数
+main() {
+    # 执行安装
+    main_install
+
+    # 清理（makeself 会自动清理临时目录）
+    exit 0
+}
+
+source ${COMMON_SHELL_PATH}
+#参数初始化
+install_file=""
+input_install_path=""
+install_for_all=n
+install_flag=n
+uninstall_flag=n
+upgrade_flag=n
+quiet_flag=n
+version_flag=n
+check_flag=n
+pylocal=y
+force_flag=n
+
+# get run package path
+run_path=`echo "$2" | cut -d"-" -f3-`
+if [ -z "${run_path}" ]; then
+    run_path=`pwd`
+else
+    # delete last "/" "/."
+    run_path=`echo "${run_path}" | sed "s/((\/)|(\/\.))*$//g"`
+    [ -z "${run_path}" ] && run_path="/"
+    if [ ! -d "${run_path}" ]; then
+        log_and_print $LEVEL_ERROR "Run package path is invalid: $run_path"
+        exit 1
+    fi
+    # 确保 run_path 是绝对路径
+    if [[ ! "${run_path}" =~ ^/ ]]; then
+        run_path="$(cd "${run_path}" && pwd)"
+    fi
+fi
+
+shift 2
+
+start_log
+
+for arg in "$@"; do
+    case "$arg" in
+    --force)
+        force_flag=y
+        ;;
+    --check)
+        check_flag=y
+        ;;
+    --quiet)
+        quiet_flag=y
+        ;;
+    --install-path=*)
+        input_install_path=`echo "$arg" | cut -d"=" -f2`
+        ;;
+    --install-for-all)
+        install_for_all=y
+        ;;
+    --full)
+        install_flag=y
+        ;;
+    --install)
+        install_flag=y
+        ;;
+    --run)
+        install_flag=y
+        ;;
+    --devel)
+        install_flag=y
+        ;;
+    --uninstall)
+        uninstall_flag=y
+        ;;
+    --upgrade)
+        upgrade_flag=y
+        ;;
+    --version)
+        version_flag=y
+        ;;
+    -*)
+        echo "Unsupported parameters: $arg"
+        ;;
+    *)
+        if [ ! -z "$arg" ]; then
+            echo "Unsupported parameters: $arg"
+        fi
+        break
+        ;;
+    esac
+done
+
+# 执行主函数
+main

msprobe/visualization/__init__.py

@@ -1,14 +1,15 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------

msprobe/visualization/builder/__init__.py

@@ -1,14 +1,15 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------

msprobe/visualization/builder/graph_builder.py

@@ -1,17 +1,18 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import re
 import copy
@@ -24,7 +25,7 @@ from msprobe.core.common.log import logger
 from msprobe.visualization.builder.msprobe_adapter import get_input_output
 from msprobe.visualization.graph.graph import Graph
 from msprobe.visualization.graph.node_op import NodeOp
-from msprobe.visualization.utils import GraphConst
+from msprobe.visualization.utils import GraphConst, update_shared_dict, update_pbar_info
 from msprobe.visualization.db_utils import node_to_db, config_to_db
 
 
@@ -36,7 +37,7 @@ class GraphBuilder:
     micro_step_dict = {}
 
     @staticmethod
-    def build(construct_path, data_path, stack_path, model_name='DefaultModel'):
+    def build(construct_path, data_path, stack_path, model_name='DefaultModel', pbar_info=None):
         """
         GraphBuilder的对外提供的构图方法
         Args:
@@ -44,6 +45,7 @@ class GraphBuilder:
             data_path: dump.json路径
             stack_path: stack.json路径
             model_name: 模型名字，依赖外部输入
+            pbar_info: 进度条对象
         Returns: Graph，代表图的数据结构
         """
         construct_dict, micro_step_dict = load_construct_json(construct_path)
@@ -58,24 +60,24 @@ class GraphBuilder:
         data_dict = dump_dict.get(GraphConst.DATA_KEY, {})
         graph = Graph(model_name, data_path=dump_dict.get('dump_data_dir', ''), dump_data=data_dict,
                       micro_step_num=micro_step_dict.get(Const.MEGATRON_MICRO_STEP_NUMBER))
-        GraphBuilder._init_nodes(graph, construct_dict, data_dict, stack_dict)
+        GraphBuilder._init_nodes(graph, construct_dict, data_dict, stack_dict, pbar_info=pbar_info)
         GraphBuilder._handle_recompute(graph)
         GraphBuilder._collect_apis_between_modules(graph)
         GraphBuilder._add_parameters_grad(graph, data_dict)
         return graph
 
     @staticmethod
-    def to_db(filename, config):
+    def to_db(filename, config, pbar_info=None):
         config.graph_n.step = config.step
         config.graph_n.rank = config.rank
         config.graph_n.compare_mode = config.compare_mode
-        node_to_db(config.graph_n, filename)
+        node_to_db(config.graph_n, filename, pbar_info)
         if config.graph_b:
             config.graph_b.data_source = GraphConst.JSON_BENCH_KEY
             config.graph_b.step = config.step
             config.graph_b.rank = config.rank
             config.graph_b.compare_mode = config.compare_mode
-            node_to_db(config.graph_b, filename)
+            node_to_db(config.graph_b, filename, pbar_info)
         config_to_db(config, filename)
 
     @staticmethod
@@ -136,8 +138,10 @@ class GraphBuilder:
         return up_node_id
 
     @staticmethod
-    def _init_nodes(graph, construct_dict, data_dict, stack_dict):
-        for subnode_id, upnode_id in construct_dict.items():
+    def _init_nodes(graph, construct_dict, data_dict, stack_dict, pbar_info=None):
+        if pbar_info:
+            update_shared_dict(pbar_info.current_stage_dict, pbar_info.task_id, 1)
+        for i, (subnode_id, upnode_id) in enumerate(construct_dict.items()):
             upnode_id = GraphBuilder._handle_backward_inplace(construct_dict, subnode_id, upnode_id) if upnode_id \
                 else GraphBuilder._handle_backward_upnode_missing(construct_dict, subnode_id, upnode_id)
             if upnode_id:
@@ -148,6 +152,9 @@ class GraphBuilder:
             node_op = NodeOp.get_node_op(subnode_id)
             GraphBuilder._create_or_get_node(graph, [data_dict, stack_dict], node_op, subnode_id, upnode)
 
+            if pbar_info:
+                update_pbar_info(pbar_info, i + 1, len(construct_dict))
+
     @staticmethod
     def _create_or_get_node(graph, data_stack_list, op, name, upnode=None):
         if name in graph.node_map:
@@ -257,13 +264,19 @@ class GraphBuilder:
         例如，graph有节点Module.a.backward.0, Module.a.backward.1, Module.a.backward.2
         则Module.a.parameters_grad添加在Module.a.backward.2内作为子节点
         """
-        prefixes = []
-        suffix = Const.SEP + Const.PARAMS_GRAD
+        prefix_suffix_map = {}
+        base_suffix = Const.SEP + Const.PARAMS_GRAD
+        pattern = re.compile(rf"{re.escape(base_suffix)}(?:\.\d+)?$")
         for node_id in data_dict.keys():
-            if node_id not in graph.node_map and node_id.endswith(suffix):
-                prefixes.append(node_id.replace(suffix, ''))
+            if node_id not in graph.node_map:
+                m = pattern.search(node_id)
+                if not m:
+                    continue
+                prefix = node_id[:m.start()]
+                suffix = node_id[m.start():]
+                prefix_suffix_map.setdefault(prefix, []).append(suffix)
 
-        max_info = {prefix: 0 for prefix in prefixes}
+        max_info = {prefix: 0 for prefix in prefix_suffix_map}
 
         for key in graph.node_map.keys():
             parts = key.split(Const.SEP)
@@ -277,12 +290,14 @@ class GraphBuilder:
             node_id = prefix + Const.SEP + Const.BACKWARD + Const.SEP + str(num)
             node = graph.get_node(node_id)
             if node:
-                parameters_grad_node_id = graph.add_node(NodeOp.module, prefix + suffix, up_node=node)
-                # 添加输入输出数据
-                node_data = data_dict.get(parameters_grad_node_id, {})
-                input_data, output_data = get_input_output(node_data, parameters_grad_node_id)
-                # 更新数据
-                graph.get_node(parameters_grad_node_id).set_input_output(input_data, output_data)
+                suffix_list = prefix_suffix_map[prefix]
+                for suffix in suffix_list:
+                    parameters_grad_node_id = graph.add_node(NodeOp.module, prefix + suffix, up_node=node)
+                    # 添加输入输出数据
+                    node_data = data_dict.get(parameters_grad_node_id, {})
+                    input_data, output_data = get_input_output(node_data, parameters_grad_node_id)
+                    # 更新数据
+                    graph.get_node(parameters_grad_node_id).set_input_output(input_data, output_data)
 
     @staticmethod
     def _handle_recompute(graph):

msprobe/visualization/builder/graph_merger.py

@@ -1,17 +1,18 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import re
 import math
@@ -20,7 +21,7 @@ from msprobe.core.common.const import Const
 from msprobe.visualization.graph.graph import Graph, BaseNode
 from msprobe.visualization.graph.node_op import NodeOp
 from msprobe.core.common.log import logger
-from msprobe.visualization.utils import GraphConst
+from msprobe.visualization.utils import GraphConst, update_shared_dict, update_pbar_info
 from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.core.common.parallel_state import get_tp_pp_default_groups
 
@@ -31,32 +32,34 @@ NORM_INFO = 'The Norm value merging method for '
 
 
 class GraphMerger:
-    def __init__(self, build_graph_results, parallel_param, is_bench=False):
-        self.strategy = self._select_strategy(build_graph_results, parallel_param, is_bench)
+    def __init__(self, build_graph_results, parallel_param, is_bench=False, pbar_info=None):
+        self.strategy = self._select_strategy(build_graph_results, parallel_param, is_bench, pbar_info)
 
     @staticmethod
-    def _select_strategy(results, param, is_bench):
+    def _select_strategy(results, param, is_bench, pbar_info):
         if param.tp == param.pp == param.rank_size == 1:
-            return NoParallelMerger(results, param, is_bench)
+            return NoParallelMerger(results, param, is_bench, pbar_info)
         elif param.tp == param.rank_size:
-            return TPMerger(results, param, is_bench)
+            return TPMerger(results, param, is_bench, pbar_info)
         elif param.pp == param.rank_size:
-            return PPMerger(results, param, is_bench) if param.vpp == 1 else VPPMerger(results, param, is_bench)
+            return PPMerger(results, param, is_bench, pbar_info) if param.vpp == 1 \
+                else VPPMerger(results, param, is_bench, pbar_info)
         elif param.pp == 1:
-            return TPMerger(results, param, is_bench)
+            return TPMerger(results, param, is_bench, pbar_info)
         elif param.tp == 1:
-            return PPMerger(results, param, is_bench) if param.vpp == 1 else VPPMerger(results, param, is_bench)
+            return PPMerger(results, param, is_bench, pbar_info) if param.vpp == 1 \
+                else VPPMerger(results, param, is_bench, pbar_info)
         elif param.tp * param.pp == param.rank_size:
-            return TPPPMerger(results, param, is_bench)
+            return TPPPMerger(results, param, is_bench, pbar_info)
         else:
-            return FullMerger(results, param, is_bench)
+            return FullMerger(results, param, is_bench, pbar_info)
 
     def merge_graph(self):
         return self.strategy.merge_graphs()
 
 
 class BaseGraphMerger:
-    def __init__(self, build_graph_results, parallel_param, is_bench):
+    def __init__(self, build_graph_results, parallel_param, is_bench, pbar_info=None):
         self.unmerged_module = [Const.CLIP_GRAD, Const.OPTIMIZER]
         self.dtype_list = Const.TORCH_INT_DTYPE + Const.TORCH_FLOAT_DTYPE + [Const.FLOAT16, Const.FLOAT32,
                                                                              Const.BFLOAT16]
@@ -64,6 +67,7 @@ class BaseGraphMerger:
         self.parallel_param = parallel_param
         self.is_bench = is_bench
         self.log_prefix = '[Bench]' if self.is_bench else '[NPU]'
+        self.pbar_info = pbar_info
         self._add_all_nodes_rank()
 
     @staticmethod
@@ -335,14 +339,19 @@ class PPMerger(BaseGraphMerger):
         return results
 
     def merge_pp_graphs(self, results):
+        if self.pbar_info:
+            update_shared_dict(self.pbar_info.current_stage_dict, self.pbar_info.task_id, 1, update_all=True)
         if not results or len(results) < 2:
             return results
         graphs = [x.graph for x in results]
         main_graph_result = results[0]
-        for main_node in main_graph_result.graph.root.subnodes:
+        total = len(main_graph_result.graph.root.subnodes)
+        for i, main_node in enumerate(main_graph_result.graph.root.subnodes):
             if main_node.op == NodeOp.module and main_node.id not in self.unmerged_module:
                 self._merge_nodes(main_graph_result.graph, main_node, graphs[1:])
                 self._sort_nodes(main_graph_result.graph, main_node)
+            if self.pbar_info:
+                update_pbar_info(self.pbar_info, i + 1, total, update_all=True)
         return [main_graph_result]
 
     def get_groups(self):
@@ -625,11 +634,14 @@ class TPMerger(BaseGraphMerger):
         return results
 
     def merge_tp_graphs(self, results, tp_merge_mapping=None):
+        if self.pbar_info:
+            update_shared_dict(self.pbar_info.current_stage_dict, self.pbar_info.task_id, 1, update_all=True)
         if not results or len(results) < 2:
             return results
         graphs = [x.graph for x in results]
         main_graph_result = results[0]
-        for main_node in main_graph_result.graph.node_map.values():
+        total = len(main_graph_result.graph.node_map.values())
+        for i, main_node in enumerate(main_graph_result.graph.node_map.values()):
             should_continue = (
                     not main_node.upnode or main_node.upnode.op != NodeOp.module or
                     main_node.upnode.id in self.unmerged_module or main_node.id.startswith(Const.DISTRIBUTED) or
@@ -647,6 +659,8 @@ class TPMerger(BaseGraphMerger):
                 merge_info_in = self._merge_params(tp_need_merge_param_in)
                 merge_info_out = self._merge_params(tp_need_merge_param_out)
                 main_node.parallel_merge_info.extend(merge_info_in + merge_info_out)
+            if self.pbar_info:
+                update_pbar_info(self.pbar_info, i + 1, total, update_all=True)
         for main_node in main_graph_result.graph.node_map.values():
             self._merge_tp_megatron_column_row_parallel(main_node, graphs[1:], tp_merge_mapping)
         return [main_graph_result]
@@ -786,9 +800,10 @@ class NoParallelMerger(BaseGraphMerger):
 
 class TPPPMerger(BaseGraphMerger):
     def merge_graphs(self):
-        tp_merger = TPMerger(self.build_graph_results, self.parallel_param, self.is_bench)
-        pp_merger = PPMerger(self.build_graph_results, self.parallel_param, self.is_bench) \
-            if self.parallel_param.vpp == 1 else VPPMerger(self.build_graph_results, self.parallel_param, self.is_bench)
+        tp_merger = TPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info)
+        pp_merger = PPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info) \
+            if self.parallel_param.vpp == 1 else \
+            VPPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info)
         pp_groups = pp_merger.get_groups()
         tp_groups = tp_merger.get_groups()
         # 进入TP+PP混合处理器，PP和TP必然大于1
@@ -809,9 +824,10 @@ class TPPPMerger(BaseGraphMerger):
 
 class FullMerger(BaseGraphMerger):
     def merge_graphs(self):
-        tp_merger = TPMerger(self.build_graph_results, self.parallel_param, self.is_bench)
-        pp_merger = PPMerger(self.build_graph_results, self.parallel_param, self.is_bench) \
-            if self.parallel_param.vpp == 1 else VPPMerger(self.build_graph_results, self.parallel_param, self.is_bench)
+        tp_merger = TPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info)
+        pp_merger = PPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info) \
+            if self.parallel_param.vpp == 1 else \
+            VPPMerger(self.build_graph_results, self.parallel_param, self.is_bench, self.pbar_info)
         pp_groups = pp_merger.get_groups()
         tp_groups = tp_merger.get_groups()
         tp_merge_mapping = {}
@@ -860,14 +876,19 @@ class VPPMerger(PPMerger):
         return Const.SEP.join(parts)
 
     def merge_pp_graphs(self, results):
+        if self.pbar_info:
+            update_shared_dict(self.pbar_info.current_stage_dict, self.pbar_info.task_id, 1, update_all=True)
         if not results or len(results) < 2:
             return results
         graphs = [x.graph for x in results]
         main_graph_result = results[0]
-        for main_node in main_graph_result.graph.root.subnodes:
+        total = len(main_graph_result.graph.root.subnodes)
+        for i, main_node in enumerate(main_graph_result.graph.root.subnodes):
             if main_node.op == NodeOp.module and main_node.id not in self.unmerged_module:
                 self._merge_nodes(main_graph_result.graph, main_node, graphs[1:])
                 self._sort_nodes(main_graph_result.graph, main_node)
+            if self.pbar_info:
+                update_pbar_info(self.pbar_info, i + 1, total, update_all=True)
         self._merge_vpp_data(main_graph_result.graph)
         self._merge_vpp_chunks(main_graph_result.graph)
         return [main_graph_result]