| 6 | |
| 7 | |
| 8 | class Lion(Optimizer1State): |
| 9 | def __init__( |
| 10 | self, |
| 11 | params, |
| 12 | lr=1e-4, |
| 13 | betas=(0.9, 0.99), |
| 14 | weight_decay=0, |
| 15 | optim_bits=32, |
| 16 | args=None, |
| 17 | min_8bit_size=4096, |
| 18 | is_paged=False, |
| 19 | ): |
| 20 | """ |
| 21 | Base Lion optimizer. |
| 22 | |
| 23 | Arguments: |
| 24 | params (`torch.tensor`): |
| 25 | The input parameters to optimize. |
| 26 | lr (`float`, defaults to 1e-4): |
| 27 | The learning rate. |
| 28 | betas (`tuple(float, float)`, defaults to (0.9, 0.999)): |
| 29 | The beta values are the decay rates of the first and second-order moment of the optimizer. |
| 30 | weight_decay (`float`, defaults to 0): |
| 31 | The weight decay value for the optimizer. |
| 32 | optim_bits (`int`, defaults to 32): |
| 33 | The number of bits of the optimizer state. |
| 34 | args (`object`, defaults to `None`): |
| 35 | An object with additional arguments. |
| 36 | min_8bit_size (`int`, defaults to 4096): |
| 37 | The minimum number of elements of the parameter tensors for 8-bit optimization. |
| 38 | is_paged (`bool`, defaults to `False`): |
| 39 | Whether the optimizer is a paged optimizer or not. |
| 40 | """ |
| 41 | super().__init__( |
| 42 | "lion", |
| 43 | params, |
| 44 | lr, |
| 45 | betas, |
| 46 | 0.0, |
| 47 | weight_decay, |
| 48 | optim_bits, |
| 49 | args, |
| 50 | min_8bit_size, |
| 51 | is_paged=is_paged, |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | class Lion8bit(Optimizer1State): |
no outgoing calls