[bugfix] fix megatron load/finetune (#5481)

Author: Jintao-Huang

swift/megatron/trainers/base.py

@@ -139,6 +139,12 @@ def _patch_load_state_dict(self):
         from megatron.training import checkpointing
         origin__load_base_checkpoint = checkpointing._load_base_checkpoint
 
+        args = get_args()
+        origin_load_state_dict = torch.nn.Module.load_state_dict
+        origin_no_load_optim = args.no_load_optim
+        origin_no_load_rng = args.no_load_rng
+        origin_finetune = args.finetune
+
         def _load_base_checkpoint(*_args, **kwargs):
             sharded_state_dict = kwargs.get('sharded_state_dict')
             if sharded_state_dict is None:
@@ -176,20 +182,17 @@ def _load_base_checkpoint(*_args, **kwargs):
                 state_dict[origin_k] = v
             return res
 
-        origin_load_state_dict = torch.nn.Module.load_state_dict
-
         def load_state_dict(self, state_dict, strict: bool = True, *args, **kwargs):
             strict = False
             return origin_load_state_dict(self, state_dict, strict, *args, **kwargs)
 
         checkpointing._load_base_checkpoint = _load_base_checkpoint
-        torch.nn.Module.load_state_dict = load_state_dict
 
-        args = get_args()
-        origin_no_load_optim = args.no_load_optim
-        origin_no_load_rng = args.no_load_rng
-        args.no_load_optim = True
-        args.no_load_rng = True
+        if args.train_type != 'full':
+            torch.nn.Module.load_state_dict = load_state_dict
+            args.no_load_optim = True
+            args.no_load_rng = True
+            args.finetune = True
 
         try:
             yield
@@ -198,6 +201,7 @@ def load_state_dict(self, state_dict, strict: bool = True, *args, **kwargs):
             torch.nn.Module.load_state_dict = origin_load_state_dict
             args.no_load_optim = origin_no_load_optim
             args.no_load_rng = origin_no_load_rng
+            args.finetune = origin_finetune
 
     def setup_model_and_optimizer(self, model_provider_func, model_type, *_args, **kwargs):
 

swift/megatron/utils/convert.py

@@ -143,6 +143,7 @@ def test_convert_precision(hf_model, mg_model, template, torch_dtype=torch.float
     'no_save_rng': True,
     'no_load_optim': True,
     'no_load_rng': True,
+    'finetune': True,
     'attention_backend': 'unfused',
 }
 
@@ -219,6 +220,8 @@ def convert_mcore2hf(args: ExportArguments) -> None:
     initialize_megatron(extra_args_provider=extra_args_provider, args_defaults=extra_args)
 
     mg_model = megatron_model_meta.model_provider()
+    if megatron_args.load is None:
+        raise ValueError('Please specify `--mcore_model`.')
     load_checkpoint([mg_model], None, None, strict=True)
     if megatron_args.adapter_load is not None:
         peft_model = prepare_mcore_model(mg_model)