add dump tool

Author: Vectorwh

msprobe_config.json

@@ -0,0 +1,16 @@
+{
+    "task": "statistics",
+    "dump_path": "./dump_path",
+    "rank": [],
+    "step": [],
+    "level": "mix",
+    "async_dump": false,
+
+    "statistics": {
+        "scope": [],
+        "list": [],
+        "tensor_list": [],
+        "data_mode": ["all"],
+        "summary_mode": "statistics"
+    }
+}

requirements/framework.txt

@@ -14,6 +14,7 @@ importlib_metadata
 jieba
 json_repair
 matplotlib
+mindstudio-probe
 modelscope>=1.23
 nltk
 numpy

swift/megatron/argument/megatron_args.py

@@ -327,6 +327,10 @@ class ExtraMegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     # qwen3_vl, qwen3_omni
     mrope_interleaved: Optional[bool] = None
 
+    # dump
+    enable_msprobe: bool = False
+    msprobe_config: str = './msprobe_config.json'
+
     @staticmethod
     def load_args_config(ckpt_dir: Optional[str]) -> Dict[str, Any]:
         res = {}

swift/megatron/trainers/base.py

@@ -503,8 +503,17 @@ def _all_reduce_metric(self,
     def train_step(self, forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, *args,
                    **kwargs):
         new_data_iterator = self._replace_data_iterator(data_iterator, model)
-        return self._origin_train_step(forward_step_func, new_data_iterator, model, optimizer, opt_param_scheduler,
+        debugger_on = self.args.enable_msprobe
+        if debugger_on:
+            from msprobe.pytorch import PrecisionDebugger
+            debugger = PrecisionDebugger(config=self.args.msprobe_config, model=model)
+            debugger.start()
+        origin_train_step_out =  self._origin_train_step(forward_step_func, new_data_iterator, model, optimizer, opt_param_scheduler,
                                        config, *args, **kwargs)
+        if debugger_on:
+            debugger.stop()
+            debugger.step()
+        return origin_train_step_out
 
     # Code borrowed from NVIDIA/Megatron-LM
     def evaluate(