From d75ec2a66f667897df0bbec1db6b91ff4ff6248f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sat, 31 Aug 2024 23:47:05 -0400 Subject: [PATCH] auto_find_batch_size is supported w deepspeed --- src/transformers/trainer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index f62184c1904..4ed8b11ca95 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -4827,10 +4827,9 @@ def create_accelerator_and_postprocess(self): 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`.") - # `auto_find_batch_size` isn't yet supported with DeepSpeed/FSDP - if (self.is_deepspeed_enabled or self.is_fsdp_enabled) and self.args.auto_find_batch_size: - wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP" - raise NotImplementedError(f"`{wrapper}` doesn't support `auto_find_batch_size`.") + # `auto_find_batch_size` isn't yet supported with FSDP + if self.is_fsdp_enabled and self.args.auto_find_batch_size: + raise NotImplementedError("`FSDP` doesn't support `auto_find_batch_size`.") def propagate_args_to_deepspeed(self, auto_find_batch_size=False): """