fix npu finetune bug

Author: yuechaowu

dnallm/configuration/configs.py

@@ -290,6 +290,7 @@ class TrainingConfig(BaseModel):
     seed: int = 42
     bf16: bool = False
     fp16: bool = False
+    gradient_checkpointing: bool = False
     load_best_model_at_end: bool = False
     metric_for_best_model: str = "eval_loss"
     report_to: list[str] = Field(

dnallm/finetune/trainer.py

@@ -193,6 +193,7 @@ def set_up_trainer(self):
         training_args.pop("hyperparameter_search", None)
         training_args.pop("use_qlora", None)
         training_args.pop("quantization_config", None)
+        self._save_safetensors = training_args.pop("save_safetensors", True)
         self.training_args = TrainingArguments(
             **training_args,
         )
@@ -379,7 +380,7 @@ def train(self, save_tokenizer: bool = True) -> dict[str, float]:
         if hasattr(self.model, "save_pretrained"):
             self.model.save_pretrained(
                 self.train_config.output_dir,
-                safe_serialization=self.trainer.args.save_safetensors,
+                safe_serialization=self._save_safetensors,
             )
         if save_tokenizer:
             self.datasets.tokenizer.save_pretrained(self.train_config.output_dir)  # type: ignore
@@ -433,7 +434,7 @@ def search(self, save_tokenizer: bool = True) -> dict[str, Any]:
         if hasattr(self.model, "save_pretrained"):
             self.model.save_pretrained(
                 self.train_config.output_dir,
-                safe_serialization=self.trainer.args.save_safetensors,
+                safe_serialization=self._save_safetensors,
             )
         if save_tokenizer:
             self.datasets.tokenizer.save_pretrained(self.train_config.output_dir)  # type: ignore