fix megatron model args (#5677)

Author: Jintao-Huang

swift/llm/argument/base_args/base_args.py

@@ -266,9 +266,6 @@ def load_args_from_ckpt(self) -> None:
             'use_chat_template',
             'response_prefix',
         ]
-        if 'megatron' in self.__class__.__name__.lower():
-            force_load_keys = []
-            load_keys.remove('use_chat_template')
         data_keys = list(f.name for f in fields(DataArguments))
         for key, old_value in old_args.items():
             if old_value is None:

swift/megatron/argument/train_args.py

@@ -3,6 +3,8 @@
 import os
 from dataclasses import dataclass
 
+import json
+
 from swift.llm import BaseArguments
 from swift.llm.argument.base_args import to_abspath
 from swift.utils import add_version_to_work_dir, get_logger, init_process_group, is_master
@@ -15,6 +17,7 @@
 @dataclass
 class MegatronTrainArguments(MegatronArguments, BaseArguments):
     add_version: bool = True
+    load_args: bool = False
 
     def init_model_args(self, tokenizer, config):
         kwargs = self.megatron_model_meta.convert_hf_config(config)
@@ -42,6 +45,16 @@ def _init_save(self):
         if is_master():
             os.makedirs(self.save, exist_ok=True)
 
+    def _init_ckpt_dir(self, adapters=None):
+        super()._init_ckpt_dir(adapters)
+        if self.ckpt_dir and self.model is None:
+            args_path = os.path.join(self.ckpt_dir, 'args.json')
+            if not os.path.exists(args_path):
+                return
+            with open(args_path, 'r', encoding='utf-8') as f:
+                old_args = json.load(f)
+            self.model = old_args.get('model')
+
     def __post_init__(self):
         self.sequence_parallel_size = self.context_parallel_size
         if self.packing: