diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index aa54f2af1bb5..f3a574ac30e8 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -978,6 +978,10 @@ def get_optimizer_cls_and_kwargs(args: TrainingArguments) -> Tuple[Any, Any]: optimizer_kwargs.update(adam_kwargs) except ImportError: raise ValueError("Trainer tried to instantiate bnb Adam8bit but bnb is not installed!") + elif args.optim == OptimizerNames.SGD: + optimizer_cls = torch.optim.SGD + elif args.optim == OptimizerNames.ADAGRAD: + optimizer_cls = torch.optim.Adagrad else: raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}") return optimizer_cls, optimizer_kwargs diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index 631fe0b3ec57..e3b760081f84 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -87,6 +87,8 @@ class OptimizerNames(ExplicitEnum): ADAMW_APEX_FUSED = "adamw_apex_fused" ADAFACTOR = "adafactor" ADAMW_BNB = "adamw_bnb_8bit" + SGD = "sgd" + ADAGRAD = "adagrad" @dataclass