required changes

Author: ved1beta

bitsandbytes/optim/adam.py

@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-
+import warnings
 from bitsandbytes.optim.optimizer import Optimizer2State
 
 
@@ -76,6 +76,8 @@ def __init__(
         betas=(0.9, 0.999),
         eps=1e-8,
         weight_decay=0,
+        amsgrad=False,
+        optim_bits=32,
         args=None,
         min_8bit_size=4096,
         percentile_clipping=100,
@@ -96,6 +98,12 @@ def __init__(
                 The epsilon value prevents division by zero in the optimizer.
             weight_decay (`float`, defaults to 0.0):
                 The weight decay value for the optimizer.
+            amsgrad (`bool`, defaults to `False`):
+                Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead.
+                Note: This parameter is not supported in Adam8bit and must be False.
+            optim_bits (`int`, defaults to 32):
+                The number of bits of the optimizer state.
+                Note: This parameter is not used in Adam8bit as it always uses 8-bit optimization.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
@@ -107,14 +115,23 @@ def __init__(
             is_paged (`bool`, defaults to `False`):
                 Whether the optimizer is a paged optimizer or not.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("Adam8bit does not support amsgrad=True")
+
+        if optim_bits != 32:
+            # We allow the default value of 32 to maintain compatibility with the function signature,
+            # but any other value is invalid since Adam8bit always uses 8-bit optimization
+            raise ValueError("Adam8bit only supports optim_bits=32 (default value for compatibility)")
+
         super().__init__(
             "adam",
             params,
             lr,
             betas,
             eps,
             weight_decay,
-            8,
+            8,  # Hardcoded to 8 bits
             args,
             min_8bit_size,
             percentile_clipping,
@@ -277,8 +294,10 @@ def __init__(
                 The weight decay value for the optimizer.
             amsgrad (`bool`, defaults to `False`):
                 Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead.
+                Note: This parameter is not supported in PagedAdam8bit and must be False.
             optim_bits (`int`, defaults to 32):
                 The number of bits of the optimizer state.
+                Note: This parameter is not used in PagedAdam8bit as it always uses 8-bit optimization.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
@@ -290,14 +309,23 @@ def __init__(
             is_paged (`bool`, defaults to `False`):
                 Whether the optimizer is a paged optimizer or not.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("PagedAdam8bit does not support amsgrad=True")
+
+        if optim_bits != 32:
+            # We allow the default value of 32 to maintain compatibility with the function signature,
+            # but any other value is invalid since PagedAdam8bit always uses 8-bit optimization
+            raise ValueError("PagedAdam8bit only supports optim_bits=32 (default value for compatibility)")
+
         super().__init__(
             "adam",
             params,
             lr,
             betas,
             eps,
             weight_decay,
-            8,
+            8,  # Hardcoded to 8 bits
             args,
             min_8bit_size,
             percentile_clipping,

bitsandbytes/optim/adamw.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 from bitsandbytes.optim.optimizer import Optimizer2State
+import warnings
 
 
 class AdamW(Optimizer2State):
@@ -98,8 +99,10 @@ def __init__(
                 The weight decay value for the optimizer.
             amsgrad (`bool`, defaults to `False`):
                 Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead.
+                Note: This parameter is not supported in AdamW8bit and must be False.
             optim_bits (`int`, defaults to 32):
                 The number of bits of the optimizer state.
+                Note: This parameter is not used in AdamW8bit as it always uses 8-bit optimization.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
@@ -111,14 +114,23 @@ def __init__(
             is_paged (`bool`, defaults to `False`):
                 Whether the optimizer is a paged optimizer or not.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("AdamW8bit does not support amsgrad=True")
+
+        if optim_bits != 32:
+            # We allow the default value of 32 to maintain compatibility with the function signature,
+            # but any other value is invalid since AdamW8bit always uses 8-bit optimization
+            raise ValueError("AdamW8bit only supports optim_bits=32 (default value for compatibility)")
+
         super().__init__(
             "adam",
             params,
             lr,
             betas,
             eps,
             weight_decay,
-            8,
+            8,  # Hardcoded to 8 bits
             args,
             min_8bit_size,
             percentile_clipping,
@@ -279,8 +291,10 @@ def __init__(
                 The weight decay value for the optimizer.
             amsgrad (`bool`, defaults to `False`):
                 Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead.
+                Note: This parameter is not supported in PagedAdamW8bit and must be False.
             optim_bits (`int`, defaults to 32):
                 The number of bits of the optimizer state.
+                Note: This parameter is not used in PagedAdamW8bit as it always uses 8-bit optimization.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
@@ -292,14 +306,23 @@ def __init__(
             is_paged (`bool`, defaults to `False`):
                 Whether the optimizer is a paged optimizer or not.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("PagedAdamW8bit does not support amsgrad=True")
+
+        if optim_bits != 32:
+            # We allow the default value of 32 to maintain compatibility with the function signature,
+            # but any other value is invalid since PagedAdamW8bit always uses 8-bit optimization
+            raise ValueError("PagedAdamW8bit only supports optim_bits=32 (default value for compatibility)")
+
         super().__init__(
             "adam",
             params,
             lr,
             betas,
             eps,
             weight_decay,
-            8,
+            8,  # Hardcoded to 8 bits
             args,
             min_8bit_size,
             percentile_clipping,