latest bnb no longer has optim_args attribute on optimizer (#3311)

* latest bnb no longer has optim_args attribute on optimizer * update the other bnb based optimizer checks
huggingface · Jan 13, 2025 · 8097343 · 8097343
1 parent c67d47a
commit 8097343
Showing 1 changed file with 20 additions and 5 deletions.
diff --git a/src/accelerate/utils/deepspeed.py b/src/accelerate/utils/deepspeed.py
@@ -50,7 +50,13 @@ def map_pytorch_optim_to_deepspeed(optimizer):
         if is_bnb_available() and not is_adaw:
             import bitsandbytes.optim as bnb_opt
 
-            is_adaw = isinstance(optimizer, (bnb_opt.AdamW, bnb_opt.AdamW32bit)) and optimizer.optim_bits == 32
+            if isinstance(optimizer, (bnb_opt.AdamW, bnb_opt.AdamW32bit)):
+                try:
+                    is_adaw = optimizer.optim_bits == 32
+                except AttributeError:
+                    is_adaw = optimizer.args.optim_bits == 32
+            else:
+                is_adaw = False
 
         if is_adaw:
             defaults["adamw_mode"] = True
@@ -64,7 +70,11 @@ def map_pytorch_optim_to_deepspeed(optimizer):
         if is_bnb_available() and not is_ada:
             import bitsandbytes.optim as bnb_opt
 
-            is_ada = isinstance(optimizer, (bnb_opt.Adagrad, bnb_opt.Adagrad32bit)) and optimizer.optim_bits == 32
+            if isinstance(optimizer, (bnb_opt.Adagrad, bnb_opt.Adagrad32bit)):
+                try:
+                    is_ada = optimizer.optim_bits == 32
+                except AttributeError:
+                    is_ada = optimizer.args.optim_bits == 32
         if is_ada:
             from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
 
@@ -74,10 +84,15 @@ def map_pytorch_optim_to_deepspeed(optimizer):
     if is_bnb_available(min_version="0.38.0") and compare_versions("deepspeed", ">=", "0.11.0"):
         from bitsandbytes.optim import Lion, Lion32bit
 
-        if isinstance(optimizer, (Lion, Lion32bit)) and optimizer.optim_bits == 32:
-            from deepspeed.ops.lion import DeepSpeedCPULion
+        if isinstance(optimizer, (Lion, Lion32bit)):
+            try:
+                is_bnb_32bits = optimizer.optim_bits == 32
+            except AttributeError:
+                is_bnb_32bits = optimizer.args.optim_bits == 32
+            if is_bnb_32bits:
+                from deepspeed.ops.lion import DeepSpeedCPULion
 
-            optimizer_class = DeepSpeedCPULion
+                optimizer_class = DeepSpeedCPULion
 
     return optimizer_class(optimizer.param_groups, **defaults)