Compare commits

...

1 Commits

Author SHA1 Message Date
Wing Lian
3202f19f52 add save_only_model arg 2024-04-10 16:09:08 -04:00
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: if self.cfg.save_safetensors is not None:
training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors 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: if self.cfg.sample_packing_eff_est:
training_arguments_kwargs[ training_arguments_kwargs[
"sample_packing_efficiency" "sample_packing_efficiency"

View File

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