Paged AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)):
(
self,
params,
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-2,
amsgrad=False,
optim_bits=32,
args=None,
min_8bit_size=4096,
)
| 178 | |
| 179 | class PagedAdamW(Optimizer2State): |
| 180 | def __init__( |
| 181 | self, |
| 182 | params, |
| 183 | lr=1e-3, |
| 184 | betas=(0.9, 0.999), |
| 185 | eps=1e-8, |
| 186 | weight_decay=1e-2, |
| 187 | amsgrad=False, |
| 188 | optim_bits=32, |
| 189 | args=None, |
| 190 | min_8bit_size=4096, |
| 191 | ): |
| 192 | """ |
| 193 | Paged AdamW optimizer. |
| 194 | |
| 195 | Arguments: |
| 196 | params (`torch.Tensor`): |
| 197 | The input parameters to optimize. |
| 198 | lr (`float`, defaults to 1e-3): |
| 199 | The learning rate. |
| 200 | betas (`tuple(float, float)`, defaults to (0.9, 0.999)): |
| 201 | The beta values are the decay rates of the first and second-order moment of the optimizer. |
| 202 | eps (`float`, defaults to 1e-8): |
| 203 | The epsilon value prevents division by zero in the optimizer. |
| 204 | weight_decay (`float`, defaults to 1e-2): |
| 205 | The weight decay value for the optimizer. |
| 206 | amsgrad (`bool`, defaults to `False`): |
| 207 | Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. |
| 208 | optim_bits (`int`, defaults to 32): |
| 209 | The number of bits of the optimizer state. |
| 210 | args (`object`, defaults to `None`): |
| 211 | An object with additional arguments. |
| 212 | min_8bit_size (`int`, defaults to 4096): |
| 213 | The minimum number of elements of the parameter tensors for 8-bit optimization. |
| 214 | """ |
| 215 | super().__init__( |
| 216 | "adam", |
| 217 | params, |
| 218 | lr, |
| 219 | betas, |
| 220 | eps, |
| 221 | weight_decay, |
| 222 | optim_bits, |
| 223 | args, |
| 224 | min_8bit_size, |
| 225 | is_paged=True, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | class PagedAdamW8bit(Optimizer2State): |