Add support for bitsandbytes

src/transformers/file_utils.py

@@ -476,6 +476,10 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
+def is_bnb_available():
+    return importlib.util.find_spec("bitsandbytes") is not None
+
+
 def is_faiss_available():
     return _faiss_available
 

src/transformers/trainer.py

@@ -847,6 +847,17 @@ def create_optimizer(self):
                 )
             else:
                 self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+                if optimizer_cls.__name__ == "Adam8bit":
+                    from torch.nn import Embedding
+
+                    import bitsandbytes
+
+                    manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+
+                    for module in self.model.modules():
+                        if isinstance(module, Embedding):
+                            manager.register_module_override(module, "weight", {"optim_bits": 32})
+                            logger.info(f"Registering bitsandbytes override for {module}")
 
         if is_sagemaker_mp_enabled():
             self.optimizer = smp.DistributedOptimizer(self.optimizer)
@@ -897,6 +908,14 @@ def get_optimizer_cls_and_kwargs(args: TrainingArguments) -> Tuple[Any, Any]:
                 optimizer_kwargs.update(adam_kwargs)
             except ImportError:
                 raise ValueError("Trainer tried to instantiate apex FusedAdam but apex is not installed!")
+        elif args.optim == OptimizerNames.ADAMW_BNB:
+            try:
+                from bitsandbytes.optim import Adam8bit
+
+                optimizer_cls = Adam8bit
+                optimizer_kwargs.update(adam_kwargs)
+            except ImportError:
+                raise ValueError("Trainer tried to instantiate bnb Adam8bit but bnb is not installed!")
         else:
             raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")
         return optimizer_cls, optimizer_kwargs

src/transformers/training_args.py

@@ -80,6 +80,7 @@ class OptimizerNames(ExplicitEnum):
     ADAMW_TORCH_XLA = "adamw_torch_xla"
     ADAMW_APEX_FUSED = "adamw_apex_fused"
     ADAFACTOR = "adafactor"
+    ADAMW_BNB = "adamw_bnb"
 
 
 @dataclass

tests/extended/test_trainer_ext.py

@@ -20,7 +20,7 @@
 from unittest.mock import patch
 
 from parameterized import parameterized
-from transformers.file_utils import is_apex_available
+from transformers.file_utils import is_apex_available, is_bnb_available
 from transformers.integrations import is_fairscale_available
 from transformers.testing_utils import (
     CaptureStderr,
@@ -71,6 +71,17 @@ def require_apex(test_case):
         return test_case
 
 
+# a candidate for testing_utils
+def require_bnb(test_case):
+    """
+    Decorator for bits and bytes (bnb) dependency
+    """
+    if not is_bnb_available():
+        return unittest.skip("test requires bnb")(test_case)
+    else:
+        return test_case
+
+
 @require_torch
 class TestTrainerExt(TestCasePlus):
     def run_seq2seq_quick(

tests/test_trainer.py

@@ -38,7 +38,7 @@
     is_torch_available,
     logging,
 )
-from transformers.file_utils import WEIGHTS_NAME, is_apex_available
+from transformers.file_utils import WEIGHTS_NAME, is_apex_available, is_bnb_available
 from transformers.testing_utils import (
     ENDPOINT_STAGING,
     PASS,
@@ -1762,6 +1762,16 @@ def hp_name(trial):
                 default_adam_kwargs,
             )
         )
+    if is_bnb_available():
+        import bitsandbytes as bnb
+
+        optim_test_params.append(
+            (
+                OptimizerNames.ADAMW_BNB,
+                bnb.optim.Adam8bit,
+                default_adam_kwargs,
+            )
+        )
 
 
 @require_torch
@@ -1787,8 +1797,8 @@ def test_optim_supported(self, name: str, expected_cls, mandatory_kwargs):
 
     def test_fused_adam(self):
         # Pretend that apex is installed and mock apex.optimizers.FusedAdam exists.
-        # Trainer.get_optimizer_cls_and_kwargs does not use FusedAdam, but only has to return a
-        # class called, so mocking apex.optimizers.FusedAdam should be fine for testing and allow
+        # Trainer.get_optimizer_cls_and_kwargs does not use FusedAdam. It only has to return the
+        # class given, so mocking apex.optimizers.FusedAdam should be fine for testing and allow
         # the test to run without requiring an apex installation.
         mock = Mock()
         modules = {
@@ -1812,6 +1822,33 @@ def test_fused_adam_no_apex(self):
             with self.assertRaises(ValueError):
                 Trainer.get_optimizer_cls_and_kwargs(args)
 
+    def test_bnb_adam8bit(self):
+        # Pretend that Bits and Bytes is installed and mock bnb.optim.Adam8bit exists.
+        # Trainer.get_optimizer_cls_and_kwargs does not use Adam8bit. It only has to return the
+        # class given, so mocking bnb.optim.Adam8bit should be fine for testing and allow
+        # the test to run without requiring a bnb installation.
+        mock = Mock()
+        modules = {
+            "bitsandbytes": mock,
+            "bitsandbytes.optim": mock.optim,
+            "bitsandbytes.optim.Adam8bit": mock.optim.Adam8bit,
+        }
+        with patch.dict("sys.modules", modules):
+            self.check_optim_and_kwargs(
+                OptimizerNames.ADAMW_BNB,
+                default_adam_kwargs,
+                mock.optim.Adam8bit,
+            )
+
+    def test_bnb_adam8bit_no_bnb(self):
+        args = TrainingArguments(optim=OptimizerNames.ADAMW_BNB, output_dir="None")
+
+        # Pretend that bnb does not exist, even if installed. By setting bnb to None, importing
+        # bnb will fail even if bnb is installed.
+        with patch.dict("sys.modules", {"bnb.optim": None}):
+            with self.assertRaises(ValueError):
+                Trainer.get_optimizer_cls_and_kwargs(args)
+
 
 @require_torch
 @require_wandb