Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions bitsandbytes/optim/adam.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,25 +88,25 @@ def __init__(
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.
.. deprecated::
Not supported in Adam8bit and must be False. This parameter exists only for
signature compatibility with :class:`Adam` and will be removed in a future version.
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.
.. deprecated::
Not used in Adam8bit. This optimizer always uses 8-bit states regardless of this
value. The parameter exists only for signature compatibility with :class:`Adam`
and will be removed in a future version.
args (`object`, defaults to `None`):
An object with additional arguments.
min_8bit_size (`int`, defaults to 4096):
The minimum number of elements of the parameter tensors for 8-bit optimization.
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__(
Expand All @@ -116,7 +116,7 @@ def __init__(
betas,
eps,
weight_decay,
8, # Hardcoded to 8 bits
8,
args,
min_8bit_size,
is_paged=is_paged,
Expand Down Expand Up @@ -258,11 +258,14 @@ def __init__(
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 PagedAdam8bit and must be False.
.. deprecated::
Not supported in PagedAdam8bit and must be False. This parameter exists only for
signature compatibility with :class:`Adam` and will be removed in a future version.
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.
.. deprecated::
Not used in PagedAdam8bit. This optimizer always uses 8-bit states regardless of
this value. The parameter exists only for signature compatibility with :class:`Adam`
and will be removed in a future version.
args (`object`, defaults to `None`):
An object with additional arguments.
min_8bit_size (`int`, defaults to 4096):
Expand Down
22 changes: 14 additions & 8 deletions bitsandbytes/optim/adamw.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,11 +88,14 @@ def __init__(
weight_decay (`float`, defaults to 1e-2):
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.
.. deprecated::
Not supported in AdamW8bit and must be False. This parameter exists only for
signature compatibility with :class:`AdamW` and will be removed in a future version.
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.
.. deprecated::
Not used in AdamW8bit. This optimizer always uses 8-bit states regardless of this
value. The parameter exists only for signature compatibility with :class:`AdamW`
and will be removed in a future version.
args (`object`, defaults to `None`):
An object with additional arguments.
min_8bit_size (`int`, defaults to 4096):
Expand Down Expand Up @@ -254,11 +257,14 @@ def __init__(
weight_decay (`float`, defaults to 1e-2):
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.
.. deprecated::
Not supported in PagedAdamW8bit and must be False. This parameter exists only for
signature compatibility with :class:`AdamW` and will be removed in a future version.
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.
.. deprecated::
Not used in PagedAdamW8bit. This optimizer always uses 8-bit states regardless of
this value. The parameter exists only for signature compatibility with :class:`AdamW`
and will be removed in a future version.
args (`object`, defaults to `None`):
An object with additional arguments.
min_8bit_size (`int`, defaults to 4096):
Expand Down
40 changes: 40 additions & 0 deletions tests/test_optim.py
Original file line number Diff line number Diff line change
Expand Up @@ -741,3 +741,43 @@ def test_adagrad8bit_rejects_non_8_optim_bits():
bnb.optim.Adagrad8bit(p, optim_bits=32)
# default (optim_bits=8) still constructs
bnb.optim.Adagrad8bit(p)


def test_adam8bit_rejects_amsgrad():
# amsgrad is not supported in Adam8bit; must be False (relates to #1261).
p = [torch.nn.Parameter(torch.randn(8, 8))]
with pytest.raises(ValueError):
bnb.optim.Adam8bit(p, amsgrad=True)
# default (amsgrad=False) still constructs
bnb.optim.Adam8bit(p)


def test_adam8bit_rejects_non_default_optim_bits():
# optim_bits is ignored (Adam8bit always uses 8-bit); guard invalid values (relates to #1261).
p = [torch.nn.Parameter(torch.randn(8, 8))]
with pytest.raises(ValueError):
bnb.optim.Adam8bit(p, optim_bits=8)
with pytest.raises(ValueError):
bnb.optim.Adam8bit(p, optim_bits=16)
# default (optim_bits=32) still constructs
bnb.optim.Adam8bit(p)


def test_adamw8bit_rejects_amsgrad():
# amsgrad is not supported in AdamW8bit; must be False (relates to #1261).
p = [torch.nn.Parameter(torch.randn(8, 8))]
with pytest.raises(ValueError):
bnb.optim.AdamW8bit(p, amsgrad=True)
# default (amsgrad=False) still constructs
bnb.optim.AdamW8bit(p)


def test_adamw8bit_rejects_non_default_optim_bits():
# optim_bits is ignored (AdamW8bit always uses 8-bit); guard invalid values (relates to #1261).
p = [torch.nn.Parameter(torch.randn(8, 8))]
with pytest.raises(ValueError):
bnb.optim.AdamW8bit(p, optim_bits=8)
with pytest.raises(ValueError):
bnb.optim.AdamW8bit(p, optim_bits=16)
# default (optim_bits=32) still constructs
bnb.optim.AdamW8bit(p)