o Ti*@s`dZddlZddlmZddlZddlmZdZdZdZ dZ d Z d Z ee e e e gZ d Zd Zd ZdZdZdZdZdZdZdZdZdZdZdZdZdZdZdZdZ dZ!dZ"d Z#d!Z$d"Z%d#Z&d$Z'd%d&Z(d'd(Z)d)d*Z*d+d,Z+d-d.Z,d/d0Z-d1d2Z.d3d4Z/d5d6Z0 d7d8Z1Gd9dde2Z3Gd:dde2Z4Gd;dde2Z5Gdz Implementation of learning rate schedules. Taken and modified from PyTorch v1.0.1 source https://github.com/pytorch/pytorch/blob/v1.1.0/torch/optim/lr_scheduler.py N) Optimizer)logger lr_schedule LRRangeTestOneCycleWarmupLR WarmupDecayLRWarmupCosineLRlr_range_test_min_lrlr_range_test_step_ratelr_range_test_step_sizelr_range_test_staircase edge_value mid_valuecycle_first_step_sizecycle_first_stair_countcycle_second_step_sizecycle_second_stair_countdecay_step_size cycle_min_lr cycle_max_lr decay_lr_rate cycle_min_mom cycle_max_momdecay_mom_rate warmup_min_lr warmup_max_lrwarmup_num_steps warmup_typeloglinearwarmup_min_ratio cos_min_ratiototal_num_stepscCs|dd}|jdtddd|jdtddd|jd td d d|jd td dd|jdtddd|jdtd dd|jdtddd|jdtddd|jdtddd|jdtd dd|jdtddd|jd td!d"d|jd#td$d%d|jd&dd'd(d)|jd*td+d,d|jd-td.d/d|jd0td$d1d|jd2td3d4d|jd5tdd6d|jd7td d8d|jd9ttd:d|jd;tddNzConvergence Tuningz!Convergence tuning configurationsz --lr_schedulezLR schedule for training.)typedefaulthelpz--lr_range_test_min_lrMbP?zStarting lr value.z--lr_range_test_step_rate?zscaling rate for LR range test.z--lr_range_test_step_sizeztraining steps per LR change.z--lr_range_test_staircaseFz(use staircase scaling for LR range test.z--cycle_first_step_sizez7size of first step of 1Cycle schedule (training steps).z--cycle_first_stair_countz&first stair count for 1Cycle schedule.z--cycle_second_step_sizezAsize of second step of 1Cycle schedule (default first_step_size).z--cycle_second_stair_countz'second stair count for 1Cycle schedule.z--decay_step_sizezAsize of intervals for applying post cycle decay (training steps).z--cycle_min_lrg{Gz?z1Cycle LR lower bound.z--cycle_max_lrg?z1Cycle LR upper bound.z--decay_lr_ratezpost cycle LR decay rate.z--cycle_momentum store_truez Enable 1Cycle momentum schedule.)r%actionr&z--cycle_min_mom皙?z1Cycle momentum lower bound.z--cycle_max_mom?z1Cycle momentum upper bound.z--decay_mom_ratezpost cycle momentum decay rate.z--warmup_min_lrrz!WarmupLR minimum/initial LR valuez--warmup_max_lrzWarmupLR maximum LR value.z--warmup_num_stepsz"WarmupLR step count for LR warmup.z --warmup_typez*WarmupLR increasing function during warmupz--warmup_min_ratiozCosine LR lower bound.z--cos_min_ratio)add_argument_group add_argumentstrfloatintboolWARMUP_LOG_RATE)parsergroupr9R/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.pyadd_tuning_arguments<sj r;cCs$t}t|}|\}}||fSN)argparseArgumentParserr;parse_known_args)r7 lr_sched_args unknown_argsr9r9r:parse_arguments|s rBcCt|tr|jdur|j|t<t|tr|jdur|j|t<t|tr-|jdur-|j|t<t|tr>|jdur@|j|t<dSdSdSr<) hasattrLR_RANGE_TEST_MIN_LRr LR_RANGE_TEST_STEP_RATEr LR_RANGE_TEST_STEP_SIZEr LR_RANGE_TEST_STAIRCASEr argsparamsr9r9r:override_lr_range_test_params   rLcCsVt|tr|jdur|j|t<t|tr|jdur|j|t<t|tr-|jdur-|j|t<t|tr<|jdur<|j|t<t|t rK|j durK|j |t <t|t rZ|j durZ|j |t <t|t ri|jduri|j|t <t|trx|jdurx|j|t<t|tr|jdur|j|t<t|tr|jdur|j|t<t|tr|jdur|j|t<dSdSdSr<)rDCYCLE_FIRST_STEP_SIZErCYCLE_FIRST_STAIR_COUNTrCYCLE_SECOND_STEP_SIZErCYCLE_SECOND_STAIR_COUNTrDECAY_STEP_SIZEr CYCLE_MIN_LRr CYCLE_MAX_LRr DECAY_LR_RATEr CYCLE_MIN_MOMr CYCLE_MAX_MOMrDECAY_MOM_RATErrIr9r9r:override_1cycle_paramss.          rYcCrCr<) rD WARMUP_MIN_LRr WARMUP_MAX_LRrWARMUP_NUM_STEPSr WARMUP_TYPErrIr9r9r:override_warmupLR_paramsrMr^cCs"t||t||t||dSr<)rLrYr^rIr9r9r:override_paramss  r_cCst|tr |jdurddtfS|jtvrdd|jfSi}|j|d<i|d<|jtkr9t||d|dfS|jtkrIt||d|dfSt ||d|dfS)Nz"--{} not specified on command linez{} is not supported LR scheduler$rK) rD LR_SCHEDULErformatVALID_LR_SCHEDULES LR_RANGE_TESTrL ONE_CYCLErYr^)rJconfigr9r9r:get_config_from_argss    rfcCsrd|vrdSd|vr dS|d}|d}|tvrdd|fS|tkr)|tdfS|tkr3|tdfS|tdfS)Nr$)Nz&LR schedule type not defined in configrK)Nz(LR schedule params not defined in configz{} is not a valid LR schedule)rbrarcrErdrTr[)rer lr_paramsr9r9r:get_lr_from_configs   ricCs*t||D]\}}||d<qdd|DS)NlrcSg|]}|dqSrjr9.0r8r9r9r: zupdate_lr..)zip) param_groupslrs param_grouprjr9r9r: update_lrs rucCs>t|tr|St|drt|jtr|jStdt|j)N optimizerz-{} is not a subclass of torch.optim.Optimizer) isinstancerrDrv TypeErrorrar$__name__)rvr9r9r:get_torch_optimizers rzc @seZdZdZ     d deded ed ed ed ef d dZddZ ddZ ddZ ddZ ddZ d!ddZddZddZdS)"ra:Sets the learning rate of each parameter group according to learning rate range test (LRRT) policy. The policy increases learning rate starting from a base value with a constant frequency, as detailed in the paper `A disciplined approach to neural network hyper-parameters: Part 1 `_ LRRT policy is used for finding maximum LR that trains a model without divergence, and can be used to configure the LR boundaries for Cyclic LR schedules. LRRT changes the learning rate after every batch. `step` should be called after a batch has been used for training. Args: optimizer (Optimizer): Wrapped optimizer. lr_range_test_min_lr (float or list): Initial learning rate which is the lower boundary in the range test for each parameter group. lr_range_test_step_size (int): Interval of training steps to increase learning rate. Default: 2000 lr_range_test_step_rate (float): Scaling rate for range test. Default: 1.0 lr_range_test_staircase (bool): Scale in staircase fashion, rather than continuous. Default: False. last_batch_iteration (int): The index of the last batch. This parameter is used when resuming a training job. Since `step()` should be invoked after each batch instead of after each epoch, this number represents the total number of *batches* computed, not the total number of epochs computed. When last_batch_iteration=-1, the schedule is started from the beginning. Default: -1 Example: >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> scheduler = LRRangeTest(optimizer) >>> data_loader = torch.utils.data.DataLoader(...) >>> for epoch in range(10): >>> for batch in data_loader: >>> train_batch(...) >>> scheduler.step() _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820 r'r(Fr*rvr r r r last_batch_iterationcCst||_t|tst|tr-t|t|jjkr'tdt|jjt|t||_ n |gt|jj|_ ||_ ||_ ||_ ||_ |rH|jn|j|_|dkr[t|jj|j |_dSdS)Nz(expected {} lr_range_test_min_lr, got {}r*)rzrvrwlisttuplelenrr ValueErrorramin_lr step_size step_rater| staircase_staircase_interval_continuous_interval interval_fnru_last_lr)selfrvr r r r r|r9r9r:__init__8s  zLRRangeTest.__init__cCstt|jd|jSN)mathfloorr3r|rrr9r9r:rSszLRRangeTest._staircase_intervalcCst|jd|jSr)r3r|rrr9r9r:rVsz LRRangeTest._continuous_intervalcCsd|j|Sr)rrrr9r9r: _get_increaseYszLRRangeTest._get_increasecs|fdd|jDS)Ncg|]}|qSr9r9)rnr  lr_increaser9r:ro^rpz&LRRangeTest.get_lr..)rrrr9rr:get_lr\szLRRangeTest.get_lrcCt|dddus Jd|jSzB Return last computed learning rate by current scheduler. rNzneed to call step() firstgetattrrrr9r9r: get_last_lr`zLRRangeTest.get_last_lrNcC0|dur |jd}||_t|jj||_dSrr|rurvrrrr)rbatch_iterationr9r9r:stepf zLRRangeTest.stepcC d|jiSNr|r|rr9r9r: state_dictl zLRRangeTest.state_dictcC|d|_dSrrrsdr9r9r:load_state_dictozLRRangeTest.load_state_dict)r'r{r(Fr*r<)ry __module__ __qualname____doc__rr3r4r5rrrrrrrrrr9r9r9r:rs8(   c@seZdZdZ           d(d d Zd d ZddZddZddZddZ ddZ ddZ ddZ ddZ ddZd d!Zd)d"d#Zd$d%Zd&d'ZdS)*raSets the learning rate of each parameter group according to 1Cycle learning rate policy (1CLR). 1CLR is a variation of the Cyclical Learning Rate (CLR) policy that involves one cycle followed by decay. The policy simultaneously cycles the learning rate (and momentum) between two boundaries with a constant frequency, as detailed in the paper `A disciplined approach to neural network hyper-parameters`_. 1CLR policy changes the learning rate after every batch. `step` should be called after a batch has been used for training. This implementation was adapted from the github repo: `PyTorch `_. Args: optimizer (Optimizer): Wrapped optimizer. cycle_min_lr (float or list): Initial learning rate which is the lower boundary in the cycle for each parameter group. cycle_max_lr (float or list): Upper learning rate boundaries in the cycle for each parameter group. Functionally, it defines the cycle amplitude (cycle_max_lr - cycle_min_lr). The lr at any cycle is the sum of cycle_min_lr and some scaling of the amplitude; therefore cycle_max_lr may not actually be reached depending on scaling function. decay_lr_rate(float): Decay rate for learning rate. Default: 0. cycle_first_step_size (int): Number of training iterations in the increasing half of a cycle. Default: 2000 cycle_second_step_size (int): Number of training iterations in the decreasing half of a cycle. If cycle_second_step_size is None, it is set to cycle_first_step_size. Default: None cycle_first_stair_count(int): Number of stairs in first half of cycle phase. This means lr/mom are changed in staircase fashion. Default 0, means staircase disabled. cycle_second_stair_count(int): Number of stairs in second half of cycle phase. This means lr/mom are changed in staircase fashion. Default 0, means staircase disabled. decay_step_size (int): Intervals for applying decay in decay phase. Default: 0, means no decay. cycle_momentum (bool): If ``True``, momentum is cycled inversely to learning rate between 'cycle_min_mom' and 'cycle_max_mom'. Default: True cycle_min_mom (float or list): Initial momentum which is the lower boundary in the cycle for each parameter group. Default: 0.8 cycle_max_mom (float or list): Upper momentum boundaries in the cycle for each parameter group. Functionally, it defines the cycle amplitude (cycle_max_mom - cycle_min_mom). The momentum at any cycle is the difference of cycle_max_mom and some scaling of the amplitude; therefore cycle_min_mom may not actually be reached depending on scaling function. Default: 0.9 decay_mom_rate (float): Decay rate for momentum. Default: 0. last_batch_iteration (int): The index of the last batch. This parameter is used when resuming a training job. Since `step()` should be invoked after each batch instead of after each epoch, this number represents the total number of *batches* computed, not the total number of epochs computed. When last_batch_iteration=-1, the schedule is started from the beginning. Default: -1 Example: >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> scheduler = OneCycle(optimizer, 0.0001, 0.0010) >>> data_loader = torch.utils.data.DataLoader(...) >>> for epoch in range(10): >>> for batch in data_loader: >>> train_batch(...) >>> scheduler.step() .. _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820 r+r{NrTr.r/r*cCsXt||_|||||| ||j||||| |_| r'||j| | | |||_dSr<)rzrv_initialize_cycle_initialize_lrcycle_momentum_initialize_momentumr|)rrvrrrrrrrrrrrrr|r9r9r:rs   zOneCycle.__init__cCs~t|}|dur t|n|}|||_||j|_||_|dur"|n||_||_t|jdr7d|_d|_ dSd|_d|_ dS)NrTF) r3 total_size step_ratiofirst_stair_countsecond_stair_countrrisclose skip_lr_decayskip_mom_decay)rrrrrrr9r9r:rs"    zOneCycle._initialize_cyclecCsn|gt|j|_|dkrt|j|jD]\}}||d<q|gt|j|_||_t|jdr5d|_dSdS)Nr*rjrT) rrrmin_lrsrqmax_lrsrrrr)rrvrrrr|rjr8r9r9r:rs  zOneCycle._initialize_lrc Csd|jvrt|j}td|dd|_dS||_|dfgt|j|_ |dfgt|j|_ |dkrEt |j |jD]\}}||d<q.)rrrr)rdecay_batch_iterationdecay_intervalrr9rr:_get_decay_mom0s  zOneCycle._get_decay_momcr)zCalculates the learning rate at batch index. This function is used after the cycle completes and post cycle decaying of lr/mom is enabled. This function treats `self.last_batch_iteration` as the last batch index. rcsg|]}|qSr9r9)rnrlr_decay_factorr9r:roDrpz*OneCycle._get_decay_lr..)rrrr)rrrrsr9rr: _get_decay_lr:s  zOneCycle._get_decay_lrcCs*|j|jkr |S||j|jdS)zCalculates the learning rate at batch index. This function treats `self.last_batch_iteration` as the last batch index. r)r|rrrrr9r9r:rHs zOneCycle.get_lrcCs4|jsdS|j|jkr|S||j|jdS)zCalculates the momentum at batch index. This function treats `self.last_batch_iteration` as the last batch index. Nr)rr|rrrrr9r9r:get_momPs  zOneCycle.get_momcCrrrrr9r9r:r[rzOneCycle.get_last_lrcCsb|dur |jd}||_t|jj||_|jr-|}t|jj|D] \}}||d<q$dSdS)z Updates the optimizer with the learning rate for the last batch index. `self.last_batch_iteration` is treated as the last batch index. If self.cycle_momentum is true, also updates optimizer momentum. Nrr) r|rurvrrrrrrrq)rrrrtrr9r9r:ras  z OneCycle.stepcCrrrrr9r9r:rrrzOneCycle.state_dictcCrrrrr9r9r:rurzOneCycle.load_state_dict) r+r{NrNrTr.r/r+r*r<)ryrrrrrrrrrrrrrrrrrrr9r9r9r:rss8H #       c @sxeZdZdZdddedfdededed ed ed ef d d Z ddZ ddZ dddZ ddZ ddZddZddZdS)ruIncrease the learning rate of each parameter group from min lr to max lr over warmup_num_steps steps, and then fix at max lr. Args: optimizer (Optimizer): Wrapped optimizer. warmup_min_lr (float or list): minimum learning rate. Default: 0 warmup_max_lr (float or list): maximum learning rate. Default: 0.001 warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log last_batch_iteration (int): The index of the last batch. Default: -1. Example: >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> scheduler = WarmupLR(optimizer) >>> data_loader = torch.utils.data.DataLoader(...) >>> for epoch in range(10): >>> for batch in data_loader: >>> train_batch(...) >>> scheduler.step() r+r'r)r*rvrrrrr|cCst||_||j|d|_||j|d|_ddt|j|jD|_td||_|t t hvr;t d|dt }||_ dt|j|_||_|d krZt|jj||_dSdS) Nrmax_lrcSsg|]\}}||qSr9r9)rnbigsmallr9r9r:rosz%WarmupLR.__init__..zUsing unknown warmup_type: z1. The increasing function is set to default (log)r(r*)rzrv _format_paramrrrq delta_lrsmaxrr6WARMUP_LINEAR_RATErrrrrinverse_log_warm_upr|rurrrr)rrvrrrrr|r9r9r:rs   zWarmupLR.__init__cs>|jdkr td|jS|fddt|j|jDS)NrDAttempting to get learning rate from scheduler before it has startedcsg|] \}}||qSr9r9)rnrdelta_lrgammar9r:rorz#WarmupLR.get_lr..)r|rrr _get_gammarqrrr9rr:rs  zWarmupLR.get_lrcCrrrrr9r9r:rrzWarmupLR.get_last_lrNcCrrrrr|r9r9r:rrz WarmupLR.stepcCrrrrr9r9r:rrzWarmupLR.state_dictcCrrrrr9r9r:rrzWarmupLR.load_state_dictcCsF|j|jkr!|jtkr|jt|jdS|jtkr!|j|jSdSr)r|rrr6rrrrrr9r9r:rs    zWarmupLR._get_gammacCZt|ts t|tr%t|t|jkr!tdt|j|t|t|S|gt|jSNz expected {} value for {}, got {}rwr}r~rrrrraFileNotFoundErrorrrv param_value param_namer9r9r:rzWarmupLR._format_paramr<)ryrrrr6rr3r4r2rrrrrrrrr9r9r9r:rys6   csReZdZdZdddedfdededed ed ed ed effd d Z ddZ Z S)ruIncrease the learning rate of each parameter group from min lr to max lr over warmup_num_steps steps, and then decay at linear rate over the remaining training steps. Args: optimizer (Optimizer): Wrapped optimizer. total_num_steps (int): total number of training steps warmup_min_lr (float or list): minimum learning rate. Default: 0 warmup_max_lr (float or list): maximum learning rate. Default: 0.001 warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log last_batch_iteration (int): The index of the last batch. Default: -1. Example: >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> scheduler = WarmupDecayLR(optimizer, 1000000) >>> data_loader = torch.utils.data.DataLoader(...) >>> for epoch in range(10): >>> for batch in data_loader: >>> train_batch(...) >>> scheduler.step() r+r'r)r*rvr#rrrrr|csF||_tt||||||||j|jkr!td||dSdS)N3total_num_steps {} is less than warmup_num_steps {})r#superrrrrrra)rrvr#rrrrr| __class__r9r:rs   zWarmupDecayLR.__init__cCsn|j|jkr!|jtkr|jt|jdS|jtkr!|j|jStdt |j |jt td|j |jS)Nrr+r() r|rrr6rrrrrr3r#rr9r9r:rs    zWarmupDecayLR._get_gamma) ryrrrr6rr4r3r2rr __classcell__r9r9rr:rs.c@s|eZdZdZdddedfdededed ed ed ed efd dZ ddZ dddZ ddZ ddZ ddZddZddZdS)r uIncrease the learning rate of each parameter group from min lr ratio to max lr ratio over warmup_num_steps steps, and then decay at cosine rate over the remaining training steps to min cosine ratio. Args: optimizer (Optimizer): Wrapped optimizer. total_num_steps (int): total number of training steps warmup_min_ratio (float or list): warmup start learning rate ratio. Default: 0 warmup_num_steps (int): number of steps to warm up from warmup_min_ratio to 1.0. Default: 1000 warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log cos_min_ratio (float): cosine end learning rate ratio. Default: 0.0001 last_batch_iteration (int): The index of the last batch. Default: -1. Example: >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> scheduler = WarmupCosineLR(optimizer, 1000000) >>> data_loader = torch.utils.data.DataLoader(...) >>> for epoch in range(10): >>> for batch in data_loader: >>> train_batch(...) >>> scheduler.step() r+r)g-C6?r*rvr#r!rr"rr|cCst||_||_||_||_||_||_td||_dt |j|_ |j|jkr2t d||dd|jjD|_|dkrLt|jj||_dSdS)Nrr(rcSrkrlr9rmr9r9r:ro4rpz+WarmupCosineLR.__init__..r*)rzrvr#r|r"rr!rrrrrrrrarrorg_lrsrurr)rrvr#r!rr"rr|r9r9r:rs   zWarmupCosineLR.__init__cCs|jdkr tddgS|j|jkr=|jtkr$|jt|jd}n |jt kr/|j|j}d|j }|j ||}|S|j|jd}|j |j}d|j }dt tj||d}td|j ||}|S)Nrrr+rr(r)r|rrrrr6rrrrr!r#r"cospir)rratio ratio_deltareal_last_stepreal_total_stepsr9r9r: get_lr_ratio:s"         zWarmupCosineLR.get_lr_ratioNcCrrrrr9r9r:rOrzWarmupCosineLR.stepcs6|jdkr tddgS|fdd|jDS)Nrrr+crr9r9)rnorg_lrlr_ratior9r:roZrpz)WarmupCosineLR.get_lr..)r|rrrrrr9rr:rUs  zWarmupCosineLR.get_lrcCrrrrr9r9r:r\rzWarmupCosineLR.get_last_lrcCrrrrr9r9r:rbrzWarmupCosineLR.state_dictcCrrrrr9r9r:rerzWarmupCosineLR.load_state_dictcCrrrrr9r9r:rhrzWarmupCosineLR._format_paramr<)ryrrrr6rr4r3r2rrrrrrrrr9r9r9r:r s:   )8rr= torch.optimrrdeepspeed.utilsrr`rcrd WARMUP_LRWARMUP_DECAY_LRWARMUP_COSINE_LRrbrErFrGrH EDGE_VALUE MID_VALUErNrOrPrQrRrSrTrUrVrWrXrZr[r\r]r6rWARMUP_MIN_RATIO COS_MIN_RATIOTOTAL_NUM_STEPSr;rBrLrYr^r_rfrirurzobjectrrrrr r9r9r9r:sp  @%  bZ3