8-bit paged Adam 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=0,
amsgrad=False,
optim_bits=32,
args=None,
min_8bit_size=4096,
is_paged=False,
)
| 231 | |
| 232 | class PagedAdam8bit(Optimizer2State): |
| 233 | def __init__( |
| 234 | self, |
| 235 | params, |
| 236 | lr=1e-3, |
| 237 | betas=(0.9, 0.999), |
| 238 | eps=1e-8, |
| 239 | weight_decay=0, |
| 240 | amsgrad=False, |
| 241 | optim_bits=32, |
| 242 | args=None, |
| 243 | min_8bit_size=4096, |
| 244 | is_paged=False, |
| 245 | ): |
| 246 | """ |
| 247 | 8-bit paged Adam optimizer. |
| 248 | |
| 249 | Arguments: |
| 250 | params (`torch.tensor`): |
| 251 | The input parameters to optimize. |
| 252 | lr (`float`, defaults to 1e-3): |
| 253 | The learning rate. |
| 254 | betas (`tuple(float, float)`, defaults to (0.9, 0.999)): |
| 255 | The beta values are the decay rates of the first and second-order moment of the optimizer. |
| 256 | eps (`float`, defaults to 1e-8): |
| 257 | The epsilon value prevents division by zero in the optimizer. |
| 258 | weight_decay (`float`, defaults to 0.0): |
| 259 | The weight decay value for the optimizer. |
| 260 | amsgrad (`bool`, defaults to `False`): |
| 261 | Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. |
| 262 | Note: This parameter is not supported in PagedAdam8bit and must be False. |
| 263 | optim_bits (`int`, defaults to 32): |
| 264 | The number of bits of the optimizer state. |
| 265 | Note: This parameter is not used in PagedAdam8bit as it always uses 8-bit optimization. |
| 266 | args (`object`, defaults to `None`): |
| 267 | An object with additional arguments. |
| 268 | min_8bit_size (`int`, defaults to 4096): |
| 269 | The minimum number of elements of the parameter tensors for 8-bit optimization. |
| 270 | is_paged (`bool`, defaults to `False`): |
| 271 | Whether the optimizer is a paged optimizer or not. |
| 272 | """ |
| 273 | # Validate unsupported parameters |
| 274 | if amsgrad: |
| 275 | raise ValueError("PagedAdam8bit does not support amsgrad=True") |
| 276 | |
| 277 | if optim_bits != 32: |
| 278 | # We allow the default value of 32 to maintain compatibility with the function signature, |
| 279 | # but any other value is invalid since PagedAdam8bit always uses 8-bit optimization |
| 280 | raise ValueError("PagedAdam8bit only supports optim_bits=32 (default value for compatibility)") |
| 281 | |
| 282 | super().__init__( |
| 283 | "adam", |
| 284 | params, |
| 285 | lr, |
| 286 | betas, |
| 287 | eps, |
| 288 | weight_decay, |
| 289 | 8, # Hardcoded to 8 bits |
| 290 | args, |