add save_only_model arg

This commit is contained in:
Wing Lian
2024-04-10 16:09:08 -04:00
parent 5ed29393e3
commit 3202f19f52
2 changed files with 4 additions and 0 deletions

View File

@@ -1058,6 +1058,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
if self.cfg.save_safetensors is not None:
training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors
if self.cfg.save_only_model is not None:
training_arguments_kwargs["save_only_model"] = self.cfg.save_only_model
if self.cfg.sample_packing_eff_est:
training_arguments_kwargs[
"sample_packing_efficiency"

View File

@@ -355,6 +355,7 @@ class ModelOutputConfig(BaseModel):
hub_model_id: Optional[str] = None
hub_strategy: Optional[str] = None
save_safetensors: Optional[bool] = None
save_only_model: Optional[bool] = None
class MLFlowConfig(BaseModel):