a hTm@sddlmZmZmZmZddlZddlmZddlmZm Z m Z m Z m Z m Z mZmZddgZGddde Zd d e d e d e_eeeeeeeeeeeeeeeeeeeeeeefeeeeeeed ddZeeeeejeejefeeeefdddZeeeeeeeeeeeeeeeeeeeeeeefeeeeeeed ddZeeddeeeeeeeeeeeeeeeeeeeeeeeeeefeeeeedddZdS))castOptional TYPE_CHECKINGUnionN)Tensor)_disable_dynamo_if_unsupported_get_scalar_dtype _maximize_doc _params_doc _to_scalar OptimizerParamsTTensorListList Adafactor adafactorc sxeZdZddddeeeefeeeeefeeee e d fd d Z fd d Z ddZ e dddZZS)r{Gz?皙鿩NgMbP??NF)foreachmaximize)paramslr beta2_decayepsd weight_decayrrc st|tr|dkrtdd|ks4td|d|ksJtd||ddurtd|dksttd|dd|dkstd|dd |kstd |d|kstd |t|||||||d } t|| dS) NrzTensor lr must be 1-elementrz%Learning rate should be >= 0 but is: z#beta2_decay should be <= 0 but is: rz epsilon1 should be >= 0 but is: z epsilon2 should be >= 0 but is: rz,Clipping threshold d should be >= 1 but is: z$weight_decay should be >= 0 but is: )rrrrrrr) isinstancernumel ValueErrordictsuper__init__) selfrrrrrrrrdefaults __class__D/var/www/auris/lib/python3.9/site-packages/torch/optim/_adafactor.pyr$s0   zAdafactor.__init__cs~t||jD]f}|dd|dD]L}|j|g}t|dkr*t|ds*t |d}tj |t d|d<q*qdS)Nrrrstepdtype) r# __setstate__ param_groups setdefaultstategetlentorchZ is_tensorfloattensorr )r%r1grouppZp_stateZstep_valr'r)r*r.=s     zAdafactor.__setstate__c Cs4|dD]$}|jdurqt|r,td|jjr 1: \\ &\hspace{10mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t) \cdot 1_m \\ &\hspace{10mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t) \\ &\hspace{10mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}\widehat{V}_t \leftarrow \widehat{\beta}_{2_t}\widehat{V}_{t-1}+ (1-\widehat{\beta}_{2_t}) \cdot (G_t \odot G_t) \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ &\hspace{5mm}\widehat{U}_t \leftarrow \frac{U_t}{max(1, \frac{\text{RMS}(U_t)}{d})} \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \alpha_t \widehat{U}_t \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned} For further details regarding the algorithm we refer to `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`_. z Args: a lr (float, Tensor, optional): unlike other optimizers, Adafactor does not require a learning rate, and Noam Shazeer and Mitchell Stern do not use lr at all. Deviating from the paper, this implementation uses lr for applying weight decay and as the maximum value for relative step size rho_t. Note that in the paper, a constant of 0.01 is used as the maximum value for relative step size, and so we set 0.01 as the default value. (default: 1e-2) beta2_decay (float, optional): the decay rate of beta2. beta2 standardly refers to the coefficient used for computing the running average of the gradient squared. (default: -0.8) eps (Tuple[float, float], optional): epsilon1 is the term added to the denominator of the update calculation to improve numerical stability. This use of epsilon1 deviates from the algorithm written in the paper! See note below for more details. epsilon2 is the term used to avoid having too small a weight update when applying parameter scaling. (default: (None, 1e-3)) d (float, optional): the clipping threshold, used to avoid larger-than-desired updates. weight_decay (float, optional): weight decay coefficient (default: 1e-2) foreach (bool, optional): whether foreach implementation of optimizer is used. Note that the foreach implementation uses ~ sizeof(params) more peak memory than the for-loop version due to the intermediates being a tensorlist vs just one tensor. As Adafactor is commonly used when memory is prohibitive, Adafactor will default to the slower single tensor for-loop implementation unless this flag is explicitly True. This behavior is contrary to other optimizers, which will attempt defaulting to foreach on CUDA for faster runtime. (default: None) a( .. Note:: The implementation of Adafactor subtly differs from Noam Shazeer and Mitchell Stern and implementations in some other frameworks with its use of learning rate and :math:`\epsilon_1`. Regarding the learning rate hyperparameter: Noam Shazeer and Mitchell Stern do not use lr at all, as the stated algorithm uses :math:`\rho_t` and update clipping to affect the step size. This implementation allows `lr` to influence the maximum value for :math:`\rho_t`: .. math:: \begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) \end{aligned} This differs from Noam Shazeer and Mitchell Stern, who use a constant of 0.01 as the maximum value of :math:`\rho_t` .. math:: \begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(0.01, \frac{1}{\sqrt{t}}) \end{aligned} Noam Shazeer and Mitchell Stern do not enforce an opinion on how weight decay should be computed, and so we use the learning rate as a coefficient for decoupled weight decay, similar to what is suggested in `Decoupled Weight Decay Regularization`_. Regarding the use of :math:`\epsilon_1`: The implementation attempts to replicate the presumed intention of Noam Shazeer and Mitchell Stern to use :math:`\epsilon_1` as a stabilizing term when the squared gradient becomes small. This stabilization can be written as .. math:: \begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ \end{aligned} where the row and column factors of gradient squared :math:`R_t` and :math:`C_t` are left alone, and we apply :math:`\epsilon_1` at the final calculation of the variance estimate :math:`\widehat{V}_t` and for the update :math:`U_t`. This is in contrast to Noam Shazeer and Mitchell Stern and other frameworks which apply :math:`\epsilon_1` to both row and column factors of the squared gradient, but not in the calculations after: .. math:: \begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{1^\top_n \cdot R_t} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{\sqrt{\widehat{V}_t}} \\ \end{aligned} .. _Adafactor\: Adaptive Learning Rates with Sublinear Memory Cost: https://arxiv.org/pdf/1804.04235 .. _Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101 )rrErFrGrHrIrKrLrrrrrMrNrrOc!Cs4|dur|dusJdtjr2t| ts:Jnt| } t|D]\}}|sX||n|| }||}||}||}||}| durt|jj } |d7}| }|| }t | d|d}t | | d |d|}| dkr|d| | |dkr|dur$|dus,Jdtj |ddd |d}|||tj |d dd |d }|||||}||jd dd j| d n.|dusJd ||}||||}|j| | d }||t d | d |d|} |j|| | dqBdS)N5Grad scaling should occur outside of optimizer.step()r?rCrow_var and col_var should be defined when grad is multidimensionalr9TrAZkeepdimr;)min0variance should be defined when grad is a vectorralpha)r4ZjitZ is_scriptingrr5r enumeratefinfor-ritemr]maxnormr Zmul_rAZsquare_Zdiv_sizeZlerp_meanZclamp_cloneZrsqrt_Zadd_)!rrErFrGrHrIrKrLrrrrrMrNrrOiparamr>Zstep_tr:r<r=Z step_floatZone_minus_beta2_tZrho_tr`Zrow_meanZcol_meanZ var_estimateZ grad_squaredupdateZdenomr)r)r*_single_tensor_adafactorGsV $      $rl) tensorlistsreturnc Cst|}i}|D]\\}}\}}||df}||df}t|dD]\}} | dus^Jd| dkr||vrdd|D||<tt|D]} ||| || |qqF||vrdd|D||<tt|D]} ||| || |qqFq|S) zGroups tensors by device, dtype, AND multidimensionality -- whether the tensor has multiple dims or just one dim (is a vector). This allows the foreach impl of Adafactor to assume that every group of params will either be factored or not.TFrNzgrad should not be NonecSsg|]}gqSr)r).0_r)r)r* zB_group_tensors_by_device_dtype_and_is_multidim..cSsg|]}gqSr)r)ror)r)r*rrrs)r Z"_group_tensors_by_device_and_dtypeitemsrarAranger3r@) rmgrouped_tensorsZultra_grouped_tensorsdevicer-rqZ matrix_keyZ vector_keyjr6rir)r)r*._group_tensors_by_device_dtype_and_is_multidims$     ryc+ s\t|dkrdS|dur |dus(Jdt| } t||||||g}|D]\\}}}\}}}}}}ttt|}ttt|}ttt|}| dur|dusJdt|j } t r|ddusJ|rt |}tj s |djr tj|tjdddddn t|dg}g}g}|D]J} || | |d| | |t| d| d q(fd d t||D}!| dkrt|d| | |rttt|}"ttt|}#|"ddur|#ddusJd d d |D}$t|$|$t|$dd |Dt|"|$|~$dd |D}%t|%|%t|%dd |Dt|#|%|~%dd t|"|#D}&dd |"D}'t|'| t|&|'~'nNttt|}(|(ddusJdt||})t|(|)|~)dd |(D}&t|&| | t|&t|&||&}*fdd t|!|*D}!t|*|!t||*qLdS)NrrXz2dtype is needed to compute eps1 when eps1 is unsetrcpu)rwr_rrYcs4g|],\}}t|d|d|qS)rZrYrdrercr )rpr8r)rNr)r*rrsz+_multi_tensor_adafactor..r[cSsg|]}tj|dddqS)r9Tr\r4rerpr>r)r)r*rr scSsg|]}|dqS)r9rfr~r)r)r*rr$rscSsg|]}tj|dddqSr;Tr\r}r~r)r)r*rr)scSsg|]}|dqS)r;rr~r)r)r*rr-rscSsg|]\}}||qSr)r))rpr:r<r)r)r*rr1scSsg|]}|jdddqSr)rg)rpr:r)r)r*rr5sr^cSsg|] }|qSr))rh)rpvr)r)r*rrFrscs:g|]2\}}| td|d|dqS)rrZrYr{)rpark)rr)r*rrNs)r3r ryrtrrBrr4rbrrZ _foreach_negcompiler is_compilingZis_cpuZ _foreach_add_r6r@rcr]zipZ _foreach_mul_Z _foreach_div_Z_foreach_lerp_Z_foreach_clamp_min_Z _foreach_mulZ_foreach_rsqrt_)+rrErFrGrHrIrKrLrrrrrMrNrrOrvrqr-Z is_multidimZdevice_params_Z device_grads_Zdevice_row_vars_Zdevice_col_vars_Zdevice_variances_Zdevice_state_steps_Z device_paramsZ device_gradsZdevice_state_stepsZone_minus_beta2_tsZbeta2_tsZrho_tssalphasZdevice_row_varsZdevice_col_varsZ row_meansZ col_meansZ var_estimatesZ row_var_meansZdevice_variancesZ grads_squaredZupdatesr))rrNr*_multi_tensor_adafactors                  r)Zsingle_tensor_fnF)rrErFrGrHrIrrKrLrOrrrrrMrNrc Cs^tjs$tdd|Ds$td|r.t}nt}|||||||| | | | |||||| ddS)zxFunctional API that performs Adafactor algorithm computation. See :class:`~torch.optim.Adafactor` for details. css|]}t|tjVqdS)N)rr4r)rptr)r)r* qszadafactor..z?`state_steps` argument must contain a list of singleton tensors) rrrrrMrNrrKrLrON)r4rrallr?rrl)rrErFrGrHrIrrKrLrOrrrrrMrNrfuncr)r)r*rVs6)NNNF)typingrrrrr4rZ optimizerrr r r r r rr__all__r__doc__rBr5rVrlr"rUrwr-ryrrr)r)r)r*s ( /K     Y #