File size: 7,315 Bytes
22c0749 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | # Copyright (c) Facebook, Inc. and its affiliates.
#
# 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
class LAMB(Optimizer2State):
def __init__(
self,
params,
lr=1e-3,
bias_correction=True,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=0,
amsgrad=False,
adam_w_mode=True,
optim_bits=32,
args=None,
min_8bit_size=4096,
max_unorm=1.0,
):
"""
Base LAMB optimizer.
Arguments:
params (`torch.tensor`):
The input parameters to optimize.
lr (`float`, defaults to 1e-3):
The learning rate.
bias_correction (`bool`, defaults to `True`):
Whether to apply bias correction to the first and second-order moments.
betas (`tuple(float, float)`, defaults to (0.9, 0.999)):
The beta values are the decay rates of the first and second-order moment of the optimizer.
eps (`float`, defaults to 1e-8):
The epsilon value prevents division by zero in the optimizer.
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.
adam_w_mode (`bool`, defaults to `True`):
Whether to use the AdamW variant.
optim_bits (`int`, defaults to 32):
The number of bits of the optimizer state.
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.
max_unorm (`float`, defaults to 1.0):
The maximum gradient norm.
"""
super().__init__(
"lamb",
params,
lr,
betas,
eps,
weight_decay,
optim_bits,
args,
min_8bit_size,
max_unorm=max_unorm,
)
class LAMB8bit(Optimizer2State):
def __init__(
self,
params,
lr=1e-3,
bias_correction=True,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=0,
amsgrad=False,
adam_w_mode=True,
args=None,
min_8bit_size=4096,
max_unorm=1.0,
):
"""
8-bit LAMB optimizer.
Arguments:
params (`torch.tensor`):
The input parameters to optimize.
lr (`float`, defaults to 1e-3):
The learning rate.
bias_correction (`bool`, defaults to `True`):
Whether to apply bias correction to the first and second-order moments.
betas (`tuple(float, float)`, defaults to (0.9, 0.999)):
The beta values are the decay rates of the first and second-order moment of the optimizer.
eps (`float`, defaults to 1e-8):
The epsilon value prevents division by zero in the optimizer.
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 LAMB8bit and must be False.
adam_w_mode (`bool`, defaults to `True`):
Whether to use the AdamW variant.
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.
max_unorm (`float`, defaults to 1.0):
The maximum update norm for trust-ratio clipping.
Note: This parameter is not supported in LAMB8bit and must be left at the
default 1.0. The 8-bit blockwise update does not implement update-norm
clipping; it is honored by the 32-bit LAMB / LAMB32bit optimizers.
"""
# Validate unsupported parameters
if amsgrad:
raise ValueError("LAMB8bit does not support amsgrad=True")
if max_unorm != 1.0:
# We allow the default value of 1.0 to maintain compatibility with the function
# signature, but the 8-bit blockwise update does not implement update-norm
# clipping, so any other value would be silently ignored.
raise ValueError("LAMB8bit only supports max_unorm=1.0 (default value for compatibility)")
super().__init__(
"lamb",
params,
lr,
betas,
eps,
weight_decay,
8,
args,
min_8bit_size,
max_unorm=max_unorm,
)
class LAMB32bit(Optimizer2State):
def __init__(
self,
params,
lr=1e-3,
bias_correction=True,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=0,
amsgrad=False,
adam_w_mode=True,
args=None,
min_8bit_size=4096,
max_unorm=1.0,
):
"""
32-bit LAMB optimizer.
Arguments:
params (`torch.tensor`):
The input parameters to optimize.
lr (`float`, defaults to 1e-3):
The learning rate.
bias_correction (`bool`, defaults to `True`):
Whether to apply bias correction to the first and second-order moments.
betas (`tuple(float, float)`, defaults to (0.9, 0.999)):
The beta values are the decay rates of the first and second-order moment of the optimizer.
eps (`float`, defaults to 1e-8):
The epsilon value prevents division by zero in the optimizer.
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.
adam_w_mode (`bool`, defaults to `True`):
Whether to use the AdamW variant.
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.
max_unorm (`float`, defaults to 1.0):
The maximum gradient norm.
"""
super().__init__(
"lamb",
params,
lr,
betas,
eps,
weight_decay,
32,
args,
min_8bit_size,
max_unorm=max_unorm,
)
|