Deprecates AdamW and adds --optim

src/transformers/optimization.py

@@ -15,6 +15,7 @@
 """PyTorch optimization for BERT model."""
 
 import math
+import warnings
 from typing import Callable, Iterable, Optional, Tuple, Union
 
 import torch
@@ -287,6 +288,8 @@ class AdamW(Optimizer):
             Decoupled weight decay to apply.
         correct_bias (:obj:`bool`, `optional`, defaults to `True`):
             Whether or not to correct bias in Adam (for instance, in Bert TF repository they use :obj:`False`).
+        no_deprecation_warning (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            A flag used to disable the deprecation warning (set to :obj:`True` to disable the warning).
     """
 
     def __init__(
@@ -297,7 +300,14 @@ def __init__(
         eps: float = 1e-6,
         weight_decay: float = 0.0,
         correct_bias: bool = True,
+        no_deprecation_warning: bool = False,
     ):
+        if not no_deprecation_warning:
+            warnings.warn(
+                "This implementation of AdamW is deprecated and will be removed in a future version. Use the"
+                "PyTorch implementation torch.optim.AdamW instead.",
+                FutureWarning,
+            )
         require_version("torch>=1.5.0")  # add_ with alpha
         if lr < 0.0:
             raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")

src/transformers/trainer.py

@@ -77,7 +77,7 @@
 from .modelcard import TrainingSummary
 from .modeling_utils import PreTrainedModel, unwrap_model
 from .models.auto.modeling_auto import MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
-from .optimization import Adafactor, AdamW, get_scheduler
+from .optimization import Adafactor, get_scheduler
 from .tokenization_utils_base import PreTrainedTokenizerBase
 from .trainer_callback import (
     CallbackHandler,
@@ -818,17 +818,43 @@ def create_optimizer(self):
                     "weight_decay": 0.0,
                 },
             ]
-            optimizer_cls = Adafactor if self.args.adafactor else AdamW
-            if self.args.adafactor:
+            if self.args.adafactor and self.args.optim not in {"adamw_hf", "adafactor"}:
+                raise ValueError(f"You passed the --adafactor flag and optimizer {self.args.optim}.")
+
+            optimizer_kwargs = {"lr": self.args.learning_rate}
+
+            adam_kwargs = {
+                "betas": (self.args.adam_beta1, self.args.adam_beta2),
+                "eps": self.args.adam_epsilon,
+            }
+
+            # TODO the following code is a good candidate for PEP 622 once Python 3.10 becomes the
+            #  minimum required version. See, https://www.python.org/dev/peps/pep-0622/
+            if (self.args.adafactor and self.args.optim == "adamw_hf") or self.args.optim == "adafactor":
                 optimizer_cls = Adafactor
-                optimizer_kwargs = {"scale_parameter": False, "relative_step": False}
-            else:
+                optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
+            elif self.args.optim == "adamw_hf":
+                from .optimization import AdamW
+
                 optimizer_cls = AdamW
-                optimizer_kwargs = {
-                    "betas": (self.args.adam_beta1, self.args.adam_beta2),
-                    "eps": self.args.adam_epsilon,
-                }
-            optimizer_kwargs["lr"] = self.args.learning_rate
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.args.optim == "adamw_torch":
+                from torch.optim import AdamW
+
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.args.optim == "apex_fused_adam":
+                try:
+                    from apex.optimizers import FusedAdam
+
+                    optimizer_cls = FusedAdam
+                    optimizer_kwargs.update(adam_kwargs)
+                except ImportError:
+                    raise ValueError(
+                        "Trainer attempted to instantiate apex.optimizers.FusedAdam but apex is not installed!"
+                    )
+            else:
+                raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {self.args.optim}")
+
             if self.sharded_ddp == ShardedDDPOption.SIMPLE:
                 self.optimizer = OSS(
                     params=optimizer_grouped_parameters,

src/transformers/training_args.py

@@ -330,9 +330,11 @@ class TrainingArguments:
             - :obj:`"tpu_metrics_debug"`: print debug metrics on TPU
 
             The options should be separated by whitespaces.
+        optim (:obj:`str`, `optional`, defaults to :obj:`adamw_hf`):
+            The optimizer to use: adamw_hf, adamw_torch, adafactor, or apex_fused_adam.
         adafactor (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether or not to use the :class:`~transformers.Adafactor` optimizer instead of
-            :class:`~transformers.AdamW`.
+            This argument is deprecated. Use ``--optim adafactor`` instead. Whether or not to use the
+            :class:`~transformers.Adafactor` optimizer instead of :class:`~transformers.AdamW`.
         group_by_length (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether or not to group together samples of roughly the same length in the training dataset (to minimize
             padding applied and be more efficient). Only useful if applying dynamic padding.
@@ -646,6 +648,10 @@ class TrainingArguments:
     label_smoothing_factor: float = field(
         default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."}
     )
+    optim: str = field(
+        default="adamw_hf",
+        metadata={"help": "The optimizer to use: adamw_hf, adamw_torch, adafactor, or apex_fused_adam."},
+    )
     adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
     group_by_length: bool = field(
         default=False,
@@ -807,6 +813,13 @@ def __post_init__(self):
                 )
             if not (self.sharded_ddp == "" or not self.sharded_ddp):
                 raise ValueError("sharded_ddp is not supported with bf16")
+
+        if self.adafactor:
+            warnings.warn(
+                "`adafactor` is deprecated and will be removed in version 5 of 🤗 Transformers. Use `--optim adafactor` instead",
+                FutureWarning,
+            )
+
         if (
             is_torch_available()
             and self.device.type != "cuda"