From d0c0cbfd80687126524774ba714d0e07f04dbeb7 Mon Sep 17 00:00:00 2001
From: pxp1 <958876660@qq.com>
Date: Wed, 13 Aug 2025 15:17:41 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=AE=89=E5=85=A8=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../msprobe/ccsrc/if/python/MsProbeIfPython.cpp     | 11 ++++++++---
 .../msprobe/ccsrc/utils/CPythonUtils.cpp            | 13 +++++++++++--
 .../config_check/ckpt_compare/megatron_loader.py    |  2 ++
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/debug/accuracy_tools/msprobe/ccsrc/if/python/MsProbeIfPython.cpp b/debug/accuracy_tools/msprobe/ccsrc/if/python/MsProbeIfPython.cpp
index fa3f65cc5..58a60538e 100644
--- a/debug/accuracy_tools/msprobe/ccsrc/if/python/MsProbeIfPython.cpp
+++ b/debug/accuracy_tools/msprobe/ccsrc/if/python/MsProbeIfPython.cpp
@@ -58,7 +58,6 @@ PyMODINIT_FUNC PyInit__msprobe_c(void)
         Py_DECREF(m);
         return nullptr;
     }
-    Py_INCREF(precisionDebugger);
 
     PyObject* cpyAgent = MindStudioDebugger::GetCPythonAgentModule();
     if (cpyAgent == nullptr) {
@@ -71,11 +70,17 @@ PyMODINIT_FUNC PyInit__msprobe_c(void)
         Py_DECREF(m);
         return nullptr;
     }
-    Py_INCREF(cpyAgent);
 
     PyMethodDef* dumpmethods = MindStudioDebugger::GetDumpMethods();
     for (PyMethodDef* method = dumpmethods; method->ml_name != nullptr; ++method) {
-        if (PyModule_AddObject(m, method->ml_name, PyCFunction_New(method, nullptr)) < 0) {
+        PyObject* func = PyCFunction_New(method, nullptr);
+        if (func == nullptr) {
+            PyErr_SetString(PyExc_ImportError, "Failed to create dump method.");
+            Py_DECREF(m);
+            return nullptr;
+        }
+        if (PyModule_AddObject(m, method->ml_name, func) < 0) {
+            Py_DECREF(func); // 释放未被模块接管的方法对象
             PyErr_SetString(PyExc_ImportError, "Failed to bind dump method.");
             Py_DECREF(m);
             return nullptr;
diff --git a/debug/accuracy_tools/msprobe/ccsrc/utils/CPythonUtils.cpp b/debug/accuracy_tools/msprobe/ccsrc/utils/CPythonUtils.cpp
index 1122a3483..9be7813a6 100644
--- a/debug/accuracy_tools/msprobe/ccsrc/utils/CPythonUtils.cpp
+++ b/debug/accuracy_tools/msprobe/ccsrc/utils/CPythonUtils.cpp
@@ -448,9 +448,18 @@ PythonListObject& PythonListObject::Insert(int64_t pos, PythonObject& item, bool
         return *this;
     }
 
-    if (PyList_Insert(ptr, pos, item) != 0) {
+    PyObject* item_ptr = item.GetPtr();  // 需确保PythonObject有此方法，返回内部PyObject*
+    if (item_ptr == nullptr) {
+        if (!ignore) {
+            PyErr_SetString(PyExc_ValueError, "Cannot insert a null object into the list.");
+        }
+        return *this;
+    }
+
+    int insert_result = PyList_Insert(ptr, pos, item_ptr);
+    if (insert_result != 0) {
         if (ignore) {
-            PyErr_Clear();
+            PyErr_Clear();  // 清除错误状态，避免影响后续操作
         }
     }
 
diff --git a/debug/accuracy_tools/msprobe/core/config_check/ckpt_compare/megatron_loader.py b/debug/accuracy_tools/msprobe/core/config_check/ckpt_compare/megatron_loader.py
index af1c5518a..5a03b4d25 100644
--- a/debug/accuracy_tools/msprobe/core/config_check/ckpt_compare/megatron_loader.py
+++ b/debug/accuracy_tools/msprobe/core/config_check/ckpt_compare/megatron_loader.py
@@ -138,6 +138,8 @@ def _consolidate_tp_weights(weights: Dict) -> Dict:
 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]
+    if not layer_idx:
+        return 1
     num_layers_per_pipeline_stage = max(layer_idx) + 1
 
     return num_layers_per_pipeline_stage
-- 
Gitee