Paged 32-bit 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,
)
| 289 | |
| 290 | class PagedAdamW32bit(Optimizer2State): |
| 291 | def __init__( |
| 292 | self, |
| 293 | params, |
| 294 | lr=1e-3, |
| 295 | betas=(0.9, 0.999), |
| 296 | eps=1e-8, |
| 297 | weight_decay=1e-2, |
| 298 | amsgrad=False, |
| 299 | optim_bits=32, |
| 300 | args=None, |
| 301 | min_8bit_size=4096, |
| 302 | ): |
| 303 | """ |
| 304 | Paged 32-bit AdamW optimizer. |
| 305 | |
| 306 | Arguments: |
| 307 | params (`torch.Tensor`): |
| 308 | The input parameters to optimize. |
| 309 | lr (`float`, defaults to 1e-3): |
| 310 | The learning rate. |
| 311 | betas (`tuple(float, float)`, defaults to (0.9, 0.999)): |
| 312 | The beta values are the decay rates of the first and second-order moment of the optimizer. |
| 313 | eps (`float`, defaults to 1e-8): |
| 314 | The epsilon value prevents division by zero in the optimizer. |
| 315 | weight_decay (`float`, defaults to 1e-2): |
| 316 | The weight decay value for the optimizer. |
| 317 | amsgrad (`bool`, defaults to `False`): |
| 318 | Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. |
| 319 | optim_bits (`int`, defaults to 32): |
| 320 | The number of bits of the optimizer state. |
| 321 | args (`object`, defaults to `None`): |
| 322 | An object with additional arguments. |
| 323 | min_8bit_size (`int`, defaults to 4096): |
| 324 | The minimum number of elements of the parameter tensors for 8-bit optimization. |
| 325 | """ |
| 326 | super().__init__( |
| 327 | "adam", |
| 328 | params, |
| 329 | lr, |
| 330 | betas, |
| 331 | eps, |
| 332 | weight_decay, |
| 333 | 32, |
| 334 | args, |
| 335 | min_8bit_size, |
| 336 | is_paged=True, |
| 337 | ) |