Raise error when using save_only_model with load_best_model_at_end for DeepSpeed/FSDP (#28866)

* Raise error when using `save_only_model` with `load_best_model_at_end` for DeepSpeed/FSDP

* Update trainer.py
This commit is contained in:
Sourab Mangrulkar 2024-02-06 11:25:44 +05:30 committed by GitHub
parent ee2a3400f2
commit 5346db1684
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -4054,6 +4054,15 @@ class Trainer:
if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None:
self.propagate_args_to_deepspeed()
# `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end`
if (
self.args.save_only_model
and (self.is_deepspeed_enabled or self.is_fsdp_enabled)
and self.args.load_best_model_at_end
):
wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP"
raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.")
def propagate_args_to_deepspeed(self, auto_find_batch_size=False):
"""
Sets values in the deepspeed plugin based on the Trainer args