Aluode/PerceptionLabPortable
0
1# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""PyTorch optimization for BERT model."""15 16import math17import warnings18from functools import partial19from typing import Optional, Union20 21import torch22from torch.optim import Optimizer23from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau24 25from .trainer_pt_utils import LayerWiseDummyOptimizer, LayerWiseDummyScheduler26from .trainer_utils import SchedulerType27from .utils import logging28 29 30logger = logging.get_logger(__name__)31 32 33def _get_constant_lambda(_=None):34 return 135 36 37def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):38 """39 Create a schedule with a constant learning rate, using the learning rate set in optimizer.40 41 Args:42 optimizer ([`~torch.optim.Optimizer`]):43 The optimizer for which to schedule the learning rate.44 last_epoch (`int`, *optional*, defaults to -1):45 The index of the last epoch when resuming training.46 47 Return:48 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.49 """50 51 return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch)52 53 54def get_reduce_on_plateau_schedule(optimizer: Optimizer, **kwargs):55 """56 Create a schedule with a constant learning rate that decreases when a metric has stopped improving.57 58 Args:59 optimizer ([`~torch.optim.Optimizer`]):60 The optimizer for which to schedule the learning rate.61 kwargs (`dict`, *optional*):62 Extra parameters to be passed to the scheduler. See `torch.optim.lr_scheduler.ReduceLROnPlateau`63 for possible parameters.64 65 Return:66 `torch.optim.lr_scheduler.ReduceLROnPlateau` with the appropriate schedule.67 """68 69 return ReduceLROnPlateau(optimizer, **kwargs)70 71 72def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int):73 if current_step < num_warmup_steps:74 return float(current_step) / float(max(1.0, num_warmup_steps))75 return 1.076 77 78def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):79 """80 Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate81 increases linearly between 0 and the initial lr set in the optimizer.82 83 Args:84 optimizer ([`~torch.optim.Optimizer`]):85 The optimizer for which to schedule the learning rate.86 num_warmup_steps (`int`):87 The number of steps for the warmup phase.88 last_epoch (`int`, *optional*, defaults to -1):89 The index of the last epoch when resuming training.90 91 Return:92 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.93 """94 95 lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps)96 return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)97 98 99def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int):100 if current_step < num_warmup_steps:101 return float(current_step) / float(max(1, num_warmup_steps))102 return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))103 104 105def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):106 """107 Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after108 a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.109 110 Args:111 optimizer ([`~torch.optim.Optimizer`]):112 The optimizer for which to schedule the learning rate.113 num_warmup_steps (`int`):114 The number of steps for the warmup phase.115 num_training_steps (`int`):116 The total number of training steps.117 last_epoch (`int`, *optional*, defaults to -1):118 The index of the last epoch when resuming training.119 120 Return:121 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.122 """123 124 lr_lambda = partial(125 _get_linear_schedule_with_warmup_lr_lambda,126 num_warmup_steps=num_warmup_steps,127 num_training_steps=num_training_steps,128 )129 return LambdaLR(optimizer, lr_lambda, last_epoch)130 131 132def _get_cosine_schedule_with_warmup_lr_lambda(133 current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float134):135 if current_step < num_warmup_steps:136 return float(current_step) / float(max(1, num_warmup_steps))137 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))138 return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))139 140 141def get_cosine_schedule_with_warmup(142 optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1143):144 """145 Create a schedule with a learning rate that decreases following the values of the cosine function between the146 initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the147 initial lr set in the optimizer.148 149 Args:150 optimizer ([`~torch.optim.Optimizer`]):151 The optimizer for which to schedule the learning rate.152 num_warmup_steps (`int`):153 The number of steps for the warmup phase.154 num_training_steps (`int`):155 The total number of training steps.156 num_cycles (`float`, *optional*, defaults to 0.5):157 The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0158 following a half-cosine).159 last_epoch (`int`, *optional*, defaults to -1):160 The index of the last epoch when resuming training.161 162 Return:163 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.164 """165 166 lr_lambda = partial(167 _get_cosine_schedule_with_warmup_lr_lambda,168 num_warmup_steps=num_warmup_steps,169 num_training_steps=num_training_steps,170 num_cycles=num_cycles,171 )172 return LambdaLR(optimizer, lr_lambda, last_epoch)173 174 175def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(176 current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int177):178 if current_step < num_warmup_steps:179 return float(current_step) / float(max(1, num_warmup_steps))180 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))181 if progress >= 1.0:182 return 0.0183 return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))184 185 186def get_cosine_with_hard_restarts_schedule_with_warmup(187 optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1188):189 """190 Create a schedule with a learning rate that decreases following the values of the cosine function between the191 initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases192 linearly between 0 and the initial lr set in the optimizer.193 194 Args:195 optimizer ([`~torch.optim.Optimizer`]):196 The optimizer for which to schedule the learning rate.197 num_warmup_steps (`int`):198 The number of steps for the warmup phase.199 num_training_steps (`int`):200 The total number of training steps.201 num_cycles (`int`, *optional*, defaults to 1):202 The number of hard restarts to use.203 last_epoch (`int`, *optional*, defaults to -1):204 The index of the last epoch when resuming training.205 206 Return:207 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.208 """209 210 lr_lambda = partial(211 _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda,212 num_warmup_steps=num_warmup_steps,213 num_training_steps=num_training_steps,214 num_cycles=num_cycles,215 )216 return LambdaLR(optimizer, lr_lambda, last_epoch)217 218 219def _get_polynomial_decay_schedule_with_warmup_lr_lambda(220 current_step: int,221 *,222 num_warmup_steps: int,223 num_training_steps: int,224 lr_end: float,225 power: float,226 lr_init: int,227):228 if current_step < num_warmup_steps:229 return float(current_step) / float(max(1, num_warmup_steps))230 elif current_step > num_training_steps:231 return lr_end / lr_init # as LambdaLR multiplies by lr_init232 else:233 lr_range = lr_init - lr_end234 decay_steps = num_training_steps - num_warmup_steps235 pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps236 decay = lr_range * pct_remaining**power + lr_end237 return decay / lr_init # as LambdaLR multiplies by lr_init238 239 240def get_polynomial_decay_schedule_with_warmup(241 optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1242):243 """244 Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the245 optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the246 initial lr set in the optimizer.247 248 Args:249 optimizer ([`~torch.optim.Optimizer`]):250 The optimizer for which to schedule the learning rate.251 num_warmup_steps (`int`):252 The number of steps for the warmup phase.253 num_training_steps (`int`):254 The total number of training steps.255 lr_end (`float`, *optional*, defaults to 1e-7):256 The end LR.257 power (`float`, *optional*, defaults to 1.0):258 Power factor.259 last_epoch (`int`, *optional*, defaults to -1):260 The index of the last epoch when resuming training.261 262 Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT263 implementation at264 https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37265 266 Return:267 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.268 269 """270 271 lr_init = optimizer.defaults["lr"]272 if not (lr_init > lr_end):273 raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})")274 275 lr_lambda = partial(276 _get_polynomial_decay_schedule_with_warmup_lr_lambda,277 num_warmup_steps=num_warmup_steps,278 num_training_steps=num_training_steps,279 lr_end=lr_end,280 power=power,281 lr_init=lr_init,282 )283 return LambdaLR(optimizer, lr_lambda, last_epoch)284 285 286def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: Optional[int] = None):287 if current_step < num_warmup_steps:288 return float(current_step) / float(max(1, num_warmup_steps))289 shift = timescale - num_warmup_steps290 decay = 1.0 / math.sqrt((current_step + shift) / timescale)291 return decay292 293 294def get_inverse_sqrt_schedule(295 optimizer: Optimizer, num_warmup_steps: int, timescale: Optional[int] = None, last_epoch: int = -1296):297 """298 Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a299 warmup period which increases lr linearly from 0 to the initial lr set in the optimizer.300 301 Args:302 optimizer ([`~torch.optim.Optimizer`]):303 The optimizer for which to schedule the learning rate.304 num_warmup_steps (`int`):305 The number of steps for the warmup phase.306 timescale (`int`, *optional*, defaults to `num_warmup_steps`):307 Time scale.308 last_epoch (`int`, *optional*, defaults to -1):309 The index of the last epoch when resuming training.310 311 Return:312 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.313 """314 # Note: this implementation is adapted from315 # https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930316 317 if timescale is None:318 timescale = num_warmup_steps or 10_000319 320 lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale)321 return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)322 323 324def _get_cosine_schedule_with_warmup_lr_lambda(325 current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float, min_lr_rate: float = 0.0326):327 if current_step < num_warmup_steps:328 return float(current_step) / float(max(1, num_warmup_steps))329 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))330 factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))331 factor = factor * (1 - min_lr_rate) + min_lr_rate332 return max(0, factor)333 334 335def get_cosine_with_min_lr_schedule_with_warmup(336 optimizer: Optimizer,337 num_warmup_steps: int,338 num_training_steps: int,339 num_cycles: float = 0.5,340 last_epoch: int = -1,341 min_lr: Optional[float] = None,342 min_lr_rate: Optional[float] = None,343):344 """345 Create a schedule with a learning rate that decreases following the values of the cosine function between the346 initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the347 initial lr set in the optimizer.348 349 Args:350 optimizer ([`~torch.optim.Optimizer`]):351 The optimizer for which to schedule the learning rate.352 num_warmup_steps (`int`):353 The number of steps for the warmup phase.354 num_training_steps (`int`):355 The total number of training steps.356 num_cycles (`float`, *optional*, defaults to 0.5):357 The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0358 following a half-cosine).359 last_epoch (`int`, *optional*, defaults to -1):360 The index of the last epoch when resuming training.361 min_lr (`float`, *optional*):362 The minimum learning rate to reach after the cosine schedule.363 min_lr_rate (`float`, *optional*):364 The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set.365 366 Return:367 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.368 """369 370 if min_lr is not None and min_lr_rate is not None:371 raise ValueError("Only one of min_lr or min_lr_rate should be set")372 elif min_lr is not None:373 min_lr_rate = min_lr / optimizer.defaults["lr"]374 elif min_lr_rate is None:375 raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`")376 377 lr_lambda = partial(378 _get_cosine_schedule_with_warmup_lr_lambda,379 num_warmup_steps=num_warmup_steps,380 num_training_steps=num_training_steps,381 num_cycles=num_cycles,382 min_lr_rate=min_lr_rate,383 )384 return LambdaLR(optimizer, lr_lambda, last_epoch)385 386 387def _get_cosine_with_min_lr_schedule_with_warmup_lr_rate_lambda(388 current_step: int,389 *,390 num_warmup_steps: int,391 num_training_steps: int,392 num_cycles: float,393 min_lr_rate: float = 0.0,394 warmup_lr_rate: Optional[float] = None,395):396 current_step = float(current_step)397 num_warmup_steps = float(num_warmup_steps)398 num_training_steps = float(num_training_steps)399 400 if current_step < num_warmup_steps:401 if warmup_lr_rate is None:402 return (current_step + 1.0) / max(1.0, num_warmup_steps)403 else:404 warmup_lr_rate = float(warmup_lr_rate)405 return warmup_lr_rate + (1.0 - warmup_lr_rate) * (current_step) / (max(1, num_warmup_steps - 1))406 progress = (current_step - num_warmup_steps + 1.0) / (max(1.0, num_training_steps - num_warmup_steps))407 factor = 0.5 * (1.0 + math.cos(math.pi * num_cycles * 2.0 * progress))408 factor = factor * (1 - min_lr_rate) + min_lr_rate409 return max(0, factor)410 411 412def get_cosine_with_min_lr_schedule_with_warmup_lr_rate(413 optimizer: Optimizer,414 num_warmup_steps: int,415 num_training_steps: int,416 num_cycles: float = 0.5,417 last_epoch: int = -1,418 min_lr: Optional[float] = None,419 min_lr_rate: Optional[float] = None,420 warmup_lr_rate: Optional[float] = None,421):422 """423 Create a schedule with a learning rate that decreases following the values of the cosine function between the424 initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the425 initial lr set in the optimizer.426 427 Args:428 optimizer ([`~torch.optim.Optimizer`]):429 The optimizer for which to schedule the learning rate.430 num_warmup_steps (`int`):431 The number of steps for the warmup phase.432 num_training_steps (`int`):433 The total number of training steps.434 num_cycles (`float`, *optional*, defaults to 0.5):435 The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0436 following a half-cosine).437 last_epoch (`int`, *optional*, defaults to -1):438 The index of the last epoch when resuming training.439 min_lr (`float`, *optional*):440 The minimum learning rate to reach after the cosine schedule.441 min_lr_rate (`float`, *optional*):442 The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set.443 warmup_lr_rate (`float`, *optional*):444 The minimum learning rate as a ratio of the start learning rate. If not set, `warmup_lr_rate` will be treated as float(1/num_warmup_steps).445 446 Return:447 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.448 """449 450 if min_lr is not None and min_lr_rate is not None:451 raise ValueError("Only one of min_lr or min_lr_rate should be set")452 elif min_lr is not None:453 min_lr_rate = min_lr / optimizer.defaults["lr"]454 elif min_lr_rate is None:455 raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`")456 457 lr_lambda = partial(458 _get_cosine_with_min_lr_schedule_with_warmup_lr_rate_lambda,459 num_warmup_steps=num_warmup_steps,460 num_training_steps=num_training_steps,461 num_cycles=num_cycles,462 min_lr_rate=min_lr_rate,463 warmup_lr_rate=warmup_lr_rate,464 )465 return LambdaLR(optimizer, lr_lambda, last_epoch)466 467 468def _get_wsd_scheduler_lambda(469 current_step: int,470 *,471 num_warmup_steps: int,472 num_stable_steps: int,473 num_decay_steps: int,474 warmup_type: str,475 decay_type: str,476 min_lr_ratio: float,477 num_cycles: float,478):479 if current_step < num_warmup_steps:480 progress = float(current_step) / float(max(1, num_warmup_steps))481 if warmup_type == "linear":482 factor = progress483 elif warmup_type == "cosine":484 factor = 0.5 * (1.0 - math.cos(math.pi * progress))485 elif warmup_type == "1-sqrt":486 factor = 1.0 - math.sqrt(1.0 - progress)487 factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio488 return max(0.0, factor)489 490 if current_step < num_warmup_steps + num_stable_steps:491 return 1.0492 493 if current_step < num_warmup_steps + num_stable_steps + num_decay_steps:494 progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps))495 if decay_type == "linear":496 factor = 1.0 - progress497 elif decay_type == "cosine":498 factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))499 elif decay_type == "1-sqrt":500 factor = 1.0 - math.sqrt(progress)501 factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio502 return max(0.0, factor)503 return min_lr_ratio504 505 506def get_wsd_schedule(507 optimizer: Optimizer,508 num_warmup_steps: int,509 num_decay_steps: int,510 num_training_steps: Optional[int] = None,511 num_stable_steps: Optional[int] = None,512 warmup_type: str = "linear",513 decay_type: str = "cosine",514 min_lr_ratio: float = 0,515 num_cycles: float = 0.5,516 last_epoch: int = -1,517):518 """519 Create a schedule with a learning rate that has three stages:520 1. warmup: increase from min_lr_ratio times the initial learning rate to the initial learning rate following a warmup_type.521 2. stable: constant learning rate.522 3. decay: decrease from the initial learning rate to min_lr_ratio times the initial learning rate following a decay_type.523 524 Args:525 optimizer ([`~torch.optim.Optimizer`]):526 The optimizer for which to schedule the learning rate.527 num_warmup_steps (`int`):528 The number of steps for the warmup phase.529 num_decay_steps (`int`):530 The number of steps for the decay phase.531 num_training_steps (`int`, *optional*):532 The total number of training steps. This is the sum of the warmup, stable and decay steps. If `num_stable_steps` is not provided, the stable phase will be `num_training_steps - num_warmup_steps - num_decay_steps`.533 num_stable_steps (`int`, *optional*):534 The number of steps for the stable phase. Please ensure that `num_warmup_steps + num_stable_steps + num_decay_steps` equals `num_training_steps`, otherwise the other steps will default to the minimum learning rate.535 warmup_type (`str`, *optional*, defaults to "linear"):536 The type of warmup to use. Can be 'linear', 'cosine' or '1-sqrt'.537 decay_type (`str`, *optional*, defaults to "cosine"):538 The type of decay to use. Can be 'linear', 'cosine' or '1-sqrt'.539 min_lr_ratio (`float`, *optional*, defaults to 0):540 The minimum learning rate as a ratio of the initial learning rate.541 num_cycles (`float`, *optional*, defaults to 0.5):542 The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0543 following a half-cosine).544 last_epoch (`int`, *optional*, defaults to -1):545 The index of the last epoch when resuming training.546 547 Return:548 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.549 """550 551 if num_training_steps is None and num_stable_steps is None:552 raise ValueError("Either num_training_steps or num_stable_steps must be specified.")553 554 if num_training_steps is not None and num_stable_steps is not None:555 warnings.warn("Both num_training_steps and num_stable_steps are specified. num_stable_steps will be used.")556 557 if warmup_type not in ["linear", "cosine", "1-sqrt"]:558 raise ValueError(f"Unknown warmup type: {warmup_type}, expected 'linear', 'cosine' or '1-sqrt'")559 560 if decay_type not in ["linear", "cosine", "1-sqrt"]:561 raise ValueError(f"Unknown decay type: {decay_type}, expected 'linear', 'cosine' or '1-sqrt'")562 563 if num_stable_steps is None:564 num_stable_steps = num_training_steps - num_warmup_steps - num_decay_steps565 566 lr_lambda = partial(567 _get_wsd_scheduler_lambda,568 num_warmup_steps=num_warmup_steps,569 num_stable_steps=num_stable_steps,570 num_decay_steps=num_decay_steps,571 warmup_type=warmup_type,572 decay_type=decay_type,573 min_lr_ratio=min_lr_ratio,574 num_cycles=num_cycles,575 )576 return LambdaLR(optimizer, lr_lambda, last_epoch)577 578 579TYPE_TO_SCHEDULER_FUNCTION = {580 SchedulerType.LINEAR: get_linear_schedule_with_warmup,581 SchedulerType.COSINE: get_cosine_schedule_with_warmup,582 SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,583 SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,584 SchedulerType.CONSTANT: get_constant_schedule,585 SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,586 SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule,587 SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule,588 SchedulerType.COSINE_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup,589 SchedulerType.COSINE_WARMUP_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup_lr_rate,590 SchedulerType.WARMUP_STABLE_DECAY: get_wsd_schedule,591}592 593 594def get_scheduler(595 name: Union[str, SchedulerType],596 optimizer: Optimizer,597 num_warmup_steps: Optional[int] = None,598 num_training_steps: Optional[int] = None,599 scheduler_specific_kwargs: Optional[dict] = None,600):601 """602 Unified API to get any scheduler from its name.603 604 Args:605 name (`str` or `SchedulerType`):606 The name of the scheduler to use.607 optimizer (`torch.optim.Optimizer`):608 The optimizer that will be used during training.609 num_warmup_steps (`int`, *optional*):610 The number of warmup steps to do. This is not required by all schedulers (hence the argument being611 optional), the function will raise an error if it's unset and the scheduler type requires it.612 num_training_steps (`int``, *optional*):613 The number of training steps to do. This is not required by all schedulers (hence the argument being614 optional), the function will raise an error if it's unset and the scheduler type requires it.615 scheduler_specific_kwargs (`dict`, *optional*):616 Extra parameters for schedulers such as cosine with restarts. Mismatched scheduler types and scheduler617 parameters will cause the scheduler function to raise a TypeError.618 """619 name = SchedulerType(name)620 schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]621 622 # If a `LayerWiseDummyOptimizer` is passed we extract the optimizer dict and623 # recursively call `get_scheduler` to get the proper schedulers on each parameter624 if optimizer is not None and isinstance(optimizer, LayerWiseDummyOptimizer):625 optimizer_dict = optimizer.optimizer_dict626 scheduler_dict = {}627 628 for param in optimizer_dict:629 scheduler_dict[param] = get_scheduler(630 name,631 optimizer=optimizer_dict[param],632 num_warmup_steps=num_warmup_steps,633 num_training_steps=num_training_steps,634 scheduler_specific_kwargs=scheduler_specific_kwargs,635 )636 637 def scheduler_hook(param):638 # Since the optimizer hook has been already attached we only need to639 # attach the scheduler hook, the gradients have been zeroed here640 scheduler_dict[param].step()641 642 for param in optimizer_dict:643 if param.requires_grad:644 param.register_post_accumulate_grad_hook(scheduler_hook)645 646 return LayerWiseDummyScheduler(optimizer_dict=optimizer_dict, lr=optimizer.defaults["lr"])647 648 if name == SchedulerType.CONSTANT:649 return schedule_func(optimizer)650 651 if scheduler_specific_kwargs is None:652 scheduler_specific_kwargs = {}653 654 if name == SchedulerType.REDUCE_ON_PLATEAU:655 return schedule_func(optimizer, **scheduler_specific_kwargs)656 657 # All other schedulers require `num_warmup_steps`658 if num_warmup_steps is None:659 raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")660 661 if name == SchedulerType.CONSTANT_WITH_WARMUP:662 return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)663 664 if name == SchedulerType.INVERSE_SQRT:665 return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)666 667 # wsd scheduler requires either num_training_steps or num_stable_steps668 if name == SchedulerType.WARMUP_STABLE_DECAY:669 return schedule_func(670 optimizer,671 num_warmup_steps=num_warmup_steps,672 num_training_steps=num_training_steps,673 **scheduler_specific_kwargs,674 )675 676 # All other schedulers require `num_training_steps`677 if num_training_steps is None:678 raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")679 680 return schedule_func(681 optimizer,682 num_warmup_steps=num_warmup_steps,683 num_training_steps=num_training_steps,684 **scheduler_specific_kwargs,685 )686 687 688class Adafactor(Optimizer):689 """690 AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:691 https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py692 693 Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://huggingface.co/papers/1804.04235 Note that694 this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and695 `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and696 `relative_step=False`.697 698 Arguments:699 params (`Iterable[nn.parameter.Parameter]`):700 Iterable of parameters to optimize or dictionaries defining parameter groups.701 lr (`float`, *optional*):702 The external learning rate.703 eps (`tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):704 Regularization constants for square gradient and parameter scale respectively705 clip_threshold (`float`, *optional*, defaults to 1.0):706 Threshold of root mean square of final gradient update707 decay_rate (`float`, *optional*, defaults to -0.8):708 Coefficient used to compute running averages of square709 beta1 (`float`, *optional*):710 Coefficient used for computing running averages of gradient711 weight_decay (`float`, *optional*, defaults to 0.0):712 Weight decay (L2 penalty)713 scale_parameter (`bool`, *optional*, defaults to `True`):714 If True, learning rate is scaled by root mean square715 relative_step (`bool`, *optional*, defaults to `True`):716 If True, time-dependent learning rate is computed instead of external learning rate717 warmup_init (`bool`, *optional*, defaults to `False`):718 Time-dependent learning rate computation depends on whether warm-up initialization is being used719 720 This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.721 722 Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):723 724 - Training without LR warmup or clip_threshold is not recommended.725 726 - use scheduled LR warm-up to fixed LR727 - use clip_threshold=1.0 (https://huggingface.co/papers/1804.04235)728 - Disable relative updates729 - Use scale_parameter=False730 - Additional optimizer operations like gradient clipping should not be used alongside Adafactor731 732 Example:733 734 ```python735 Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)736 ```737 738 Others reported the following combination to work well:739 740 ```python741 Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)742 ```743 744 When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]745 scheduler as following:746 747 ```python748 from transformers.optimization import Adafactor, AdafactorSchedule749 750 optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)751 lr_scheduler = AdafactorSchedule(optimizer)752 trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))753 ```754 755 Usage:756 757 ```python758 # replace AdamW with Adafactor759 optimizer = Adafactor(760 model.parameters(),761 lr=1e-3,762 eps=(1e-30, 1e-3),763 clip_threshold=1.0,764 decay_rate=-0.8,765 beta1=None,766 weight_decay=0.0,767 relative_step=False,768 scale_parameter=False,769 warmup_init=False,770 )771 ```"""772 773 def __init__(774 self,775 params,776 lr=None,777 eps=(1e-30, 1e-3),778 clip_threshold=1.0,779 decay_rate=-0.8,780 beta1=None,781 weight_decay=0.0,782 scale_parameter=True,783 relative_step=True,784 warmup_init=False,785 ):786 if lr is not None and relative_step:787 raise ValueError("Cannot combine manual `lr` and `relative_step=True` options")788 if warmup_init and not relative_step:789 raise ValueError("`warmup_init=True` requires `relative_step=True`")790 791 defaults = {792 "lr": lr,793 "eps": eps,794 "clip_threshold": clip_threshold,795 "decay_rate": decay_rate,796 "beta1": beta1,797 "weight_decay": weight_decay,798 "scale_parameter": scale_parameter,799 "relative_step": relative_step,800 "warmup_init": warmup_init,801 }802 super().__init__(params, defaults)803 804 @staticmethod805 def _get_lr(param_group, param_state):806 rel_step_sz = param_group["lr"]807 if param_group["relative_step"]:808 min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2809 rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))810 param_scale = 1.0811 if param_group["scale_parameter"]:812 param_scale = max(param_group["eps"][1], param_state["RMS"])813 return param_scale * rel_step_sz814 815 @staticmethod816 def _get_options(param_group, param_shape):817 factored = len(param_shape) >= 2818 use_first_moment = param_group["beta1"] is not None819 return factored, use_first_moment820 821 @staticmethod822 def _rms(tensor):823 return tensor.norm(2) / (tensor.numel() ** 0.5)824 825 @staticmethod826 def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):827 # copy from fairseq's adafactor implementation:828 # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505829 r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)830 c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()831 return torch.mul(r_factor, c_factor)832 833 @torch.no_grad()834 def step(self, closure=None):835 """836 Performs a single optimization step837 838 Arguments:839 closure (callable, optional): A closure that reevaluates the model840 and returns the loss.841 """842 loss = None843 if closure is not None:844 loss = closure()845 846 for group in self.param_groups:847 for p in group["params"]:848 if p.grad is None:849 continue850 grad = p.grad851 if grad.dtype in {torch.float16, torch.bfloat16}:852 grad = grad.float()853 if grad.is_sparse:854 raise RuntimeError("Adafactor does not support sparse gradients.")855 856 state = self.state[p]857 grad_shape = grad.shape858 859 factored, use_first_moment = self._get_options(group, grad_shape)860 # State Initialization861 if len(state) == 0:862 state["step"] = 0863 864 if use_first_moment:865 # Exponential moving average of gradient values866 state["exp_avg"] = torch.zeros_like(grad)867 if factored:868 state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)869 state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)870 else:871 state["exp_avg_sq"] = torch.zeros_like(grad)872 873 state["RMS"] = 0874 else:875 if use_first_moment:876 state["exp_avg"] = state["exp_avg"].to(grad)877 if factored:878 state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)879 state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)880 else:881 state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)882 883 p_data_fp32 = p884 if p.dtype in {torch.float16, torch.bfloat16}:885 p_data_fp32 = p_data_fp32.float()886 887 state["step"] += 1888 state["RMS"] = self._rms(p_data_fp32)889 lr = self._get_lr(group, state)890 891 beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])892 update = (grad**2) + group["eps"][0]893 if factored:894 exp_avg_sq_row = state["exp_avg_sq_row"]895 exp_avg_sq_col = state["exp_avg_sq_col"]896 897 exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))898 exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))899 900 # Approximation of exponential moving average of square of gradient901 update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)902 update.mul_(grad)903 else:904 exp_avg_sq = state["exp_avg_sq"]905 906 exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))907 update = exp_avg_sq.rsqrt().mul_(grad)908 909 update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))910 update.mul_(lr)911 912 if use_first_moment:913 exp_avg = state["exp_avg"]914 exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))915 update = exp_avg916 917 if group["weight_decay"] != 0:918 p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))919 920 p_data_fp32.add_(-update)921 922 if p.dtype in {torch.float16, torch.bfloat16}:923 p.copy_(p_data_fp32)924 925 return loss926 927 928class AdafactorSchedule(LambdaLR):929 """930 Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g.,931 for logging), this class creates a proxy object that retrieves the current lr values from the optimizer.932 933 It returns `initial_lr` during startup and the actual `lr` during stepping.934 """935 936 def __init__(self, optimizer, initial_lr=0.0):937 def lr_lambda(_):938 return initial_lr939 940 for group in optimizer.param_groups:941 group["initial_lr"] = initial_lr942 super().__init__(optimizer, lr_lambda)943 for group in optimizer.param_groups:944 del group["initial_lr"]945 946 def get_lr(self):947 opt = self.optimizer948 lrs = [949 opt._get_lr(group, opt.state[group["params"][0]])950 for group in opt.param_groups951 if group["params"][0].grad is not None952 ]953 if len(lrs) == 0:954 lrs = self.base_lrs # if called before stepping955 return lrs956 957 958def get_adafactor_schedule(optimizer, initial_lr=0.0):959 """960 Get a proxy schedule for [`~optimization.Adafactor`]961 962 Args:963 optimizer ([`~torch.optim.Optimizer`]):964 The optimizer for which to schedule the learning rate.965 initial_lr (`float`, *optional*, defaults to 0.0):966 Initial lr967 968 Return:969 [`~optimization.Adafactor`] proxy schedule object.970 971 972 """973 return AdafactorSchedule(optimizer, initial_lr)974