move prepare_model

swift/llm/train/pt.py

@@ -9,12 +9,9 @@ class SwiftPt(SwiftSft):
     args_class = TrainArguments
     args: args_class
 
-    def _prepare_train(self):
-        self.template.loss_scale = 'all'
-        super()._prepare_train()
-
     def _prepare_template(self, use_chat_template: bool) -> None:
         super()._prepare_template(use_chat_template=False)
+        self.template.loss_scale = 'all'
 
 
 def pt_main(args: Union[List[str], TrainArguments, None] = None):

swift/llm/train/rlhf.py

@@ -20,21 +20,16 @@ def _prepare_model_tokenizer(self):
 
         super()._prepare_model_tokenizer()
 
-    def _prepare_train(self):
+    def _prepare_template(self, use_chat_template: bool) -> None:
         args = self.args
+        super()._prepare_template(use_chat_template=use_chat_template)
         mode = 'kto' if args.rlhf_type == 'kto' else 'rlhf'
         self.template.set_mode(mode)
 
         if args.rlhf_type != 'orpo' or args.model_meta.is_multimodal:
             # Avoid padding labels during the model's forward pass in multimodal models.
             self.template.loss_scale = 'last_round'
 
-        if self.model.model_meta.is_multimodal:
-            models = [self.model]
-            if self.ref_model:
-                models.append(self.ref_model)
-            self.template.register_post_encode_hook(models)
-
     def _get_dataset(self):
         args = self.args
         train_dataset, val_dataset = super()._get_dataset()

swift/llm/train/sft.py

@@ -34,18 +34,6 @@ def __init__(self, args: Union[List[str], TrainArguments, None] = None) -> None:
         self._prepare_model_tokenizer()
         self._prepare_template(True)
         self._prepare_callbacks()
-        self.model = prepare_model(self.args, self.model)
-        logger.info(f'model: {self.model}')
-        model_parameter_info = get_model_parameter_info(self.model)
-        self.train_msg['model_parameter_info'] = model_parameter_info
-        logger.info(f'model_parameter_info: {model_parameter_info}')
-
-        self._prepare_train()
-
-    def _prepare_train(self):
-        self.template.set_mode('train')
-        if self.model.model_meta.is_multimodal:
-            self.template.register_post_encode_hook([self.model])
 
     def _prepare_gradient_checkpointing(self):
         args = self.args
@@ -118,6 +106,7 @@ def _prepare_template(self, use_chat_template: bool) -> None:
         logger.info(f'default_system: {template.template_meta.default_system}')
         if template.use_model:
             template.model = self.model
+        template.set_mode('train')
         self.template = template
 
     def _get_dataset(self):
@@ -152,9 +141,14 @@ def run(self):
 
         train_dataset, val_dataset = self._get_dataset()
         train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset)
+        # Some tuners require train_dataset for preparation: LoRA-GA
+        self.model = prepare_model(self.args, self.model)
+        logger.info(f'model: {self.model}')
+        model_parameter_info = get_model_parameter_info(self.model)
+        self.train_msg['model_parameter_info'] = model_parameter_info
+        logger.info(f'model_parameter_info: {model_parameter_info}')
 
         data_collator = self._get_data_collator()
-
         optimizers = self._get_optimizers(train_dataset)
 
         trainer_cls = TrainerFactory.get_trainer_cls(args)
@@ -221,7 +215,6 @@ def _save_trainer_state(self, trainer):
     def train(self, trainer):
         logging_path = os.path.join(trainer.args.output_dir, 'logging.jsonl')
         logger.info(f'The logging file will be saved in: {logging_path}')
-        trainer.model_accepts_loss_kwargs = True  # fix transformers>=4.46.2
         trainer.train(trainer.args.resume_from_checkpoint)
 
         return self._save_trainer_state(trainer)

swift/trainers/mixin.py

@@ -10,6 +10,7 @@
 
 import safetensors
 import torch
+import torch.nn as nn
 import transformers
 from datasets import Dataset as HfDataset
 from modelscope import check_local_model_is_latest
@@ -195,8 +196,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
         # tokenizer
         if not is_adapter:
             from swift.llm import save_checkpoint
-            additional_saved_files = self.model.model_meta.additional_saved_files if hasattr(self.model,
-                                                                                             'model_meta') else []
+            additional_saved_files = self.model.model_meta.additional_saved_files
             save_checkpoint(None, self.template.processor, output_dir, additional_saved_files=additional_saved_files)
 
     def _fix_zero3_gather_all_parameters(self) -> None:
@@ -226,9 +226,19 @@ def _save_checkpoint(self, *args, **kwargs):
         return result
 
     def train(self, *args, **kwargs):
+        if self.model.model_meta.is_multimodal:
+            models = list(
+                set([
+                    v for k, v in self.__dict__.items()
+                    if isinstance(v, nn.Module) and k in {'model', 'ref_model', 'reward_model', 'value_model'}
+                ]))
+            self.template.register_post_encode_hook(models)
+            logger.info(f'Successfully registered post_encode hook: {[model.__class__.__name__ for model in models]}')
+        self.model_accepts_loss_kwargs = True  # fix transformers>=4.46.2
         self._save_initial_model(self.args.output_dir)
         with self.hub.patch_hub():
             return super().train(*args, **kwargs)
+        self.template.remove_post_encode_hook()
 
     def push_to_hub(self, *args, **kwargs):
         with self.hub.patch_hub():