Enable customized optimizer for DeepSpeed (huggingface#32049)

* transformers: enable custom optimizer for DeepSpeed * transformers: modify error message --------- Co-authored-by: datakim1201 <[email protected]>
sbucaille · Oct 7, 2024 · 55be7c4 · 55be7c4
1 parent 7bae833
commit 55be7c4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -599,11 +599,11 @@ def __init__(
                     " `Trainer`. Make sure the lines `import torch_xla.core.xla_model as xm` and"
                     " `model.to(xm.xla_device())` is performed before the optimizer creation in your script."
                 )
-        if (self.is_deepspeed_enabled or self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and (
+        if (self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and (
             self.optimizer is not None or self.lr_scheduler is not None
         ):
             raise RuntimeError(
-                "Passing `optimizers` is not allowed if Deepspeed or PyTorch FSDP is enabled. "
+                "Passing `optimizers` is not allowed if PyTorch FSDP is enabled. "
                 "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
             )
         default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)