o ZŽh|uã@sddlZddlZddlZddlmZddlmZddlmZe  e ¡Z Gdd„dej ƒZ Gdd „d ej ƒZGd d „d ej ƒZGd d „d ej ƒZGdd„dej ƒZGdd„dej ƒZd$dd„Zd$dd„Zd$dd„ZGdd„deƒZGdd„deƒZGdd„deƒZd%d d!„ZGd"d#„d#eƒZdS)&éN)Únn)ÚFunctioné)Úloggingcs>eZdZdZ         d ‡fdd„ Zd d d „Z‡ZS) ÚQuantEmbeddingaÞ Quantized version of `torch.nn.Embedding`. Adds quantization-specific arguments on top of `torch.nn.Embedding`. Args: weight_bit (`int`, *optional*, defaults to `8`): Bitwidth for the quantized weight. momentum (`float`, *optional*, defaults to `0.95`): Momentum for updating the activation quantization range. quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. Nç@Féçffffffî?c s”tƒ ¡||_||_||_||_||_||_||_t   t   ||g¡¡|_ | dt   d¡¡| dt  |j ¡¡| |_| |_| |_d|_tj|_dS)NÚweight_scaling_factoréÚweight_integerF)ÚsuperÚ__init__Znum_ÚdimÚ padding_idxÚmax_normÚ norm_typeÚscale_grad_by_freqÚsparserÚ ParameterÚtorchÚzerosÚweightÚregister_bufferÚ zeros_likeÚ weight_bitÚmomentumÚ quant_modeÚpercentile_modeÚSymmetricQuantFunctionÚapplyÚweight_function) ÚselfZnum_embeddingsZ embedding_dimrrrrrZ_weightrrr©Ú __class__©úV/var/www/auris/lib/python3.10/site-packages/transformers/models/ibert/quant_modules.pyr,s  zQuantEmbedding.__init__c Csº|jstj ||j|j|j|j|j|j ¡dfS|j}|j   ¡}|  ¡  d¡}| ¡  d¡}t|j||dƒ|_| |j|j|j|j¡|_tj ||j|j|j|j|j|j ¡}||j|jfS)Nr F)rrÚ functionalZ embeddingrrrrrrÚdataÚdetachÚminÚexpandÚmaxÚ$symmetric_linear_quantization_paramsrr r!rr ) r"ÚxZ positionsZincremental_stateÚwÚ w_transformÚw_minÚw_maxZemb_intr%r%r&ÚforwardMs<ù ö ÿù zQuantEmbedding.forward) NNrFFNrr F©NN)Ú__name__Ú __module__Ú __qualname__Ú__doc__rr3Ú __classcell__r%r%r#r&rsô!rcs>eZdZdZd ‡fdd„ Zdd„Z     d d d „Z‡ZS) ÚQuantActap Quantizes the given activation. Args: activation_bit (`int`): Bitwidth for the quantized activation. act_range_momentum (`float`, *optional*, defaults to `0.95`): Momentum for updating the activation quantization range. per_channel (`bool`, *optional*, defaults to `False`): Whether to or not use channel-wise quantization. channel_len (`int`, *optional*): Specify the channel length when set the *per_channel* True. quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. r FNcs”tƒ ¡||_||_||_||_d|_tj|_ |jsF|  dt   d¡¡|  dt   d¡¡|  dt   d¡¡|j d8_ |jd7_dStdƒ‚)NFÚx_minr Úx_maxÚact_scaling_factorgñh㈵øä>ú;per-channel mode is not currently supported for activation.)r rÚactivation_bitÚact_range_momentumrÚ per_channelÚ percentilerr Ú act_functionrrrr;r<ÚNotImplementedError)r"r?r@rAZ channel_lenrr#r%r&rƒs zQuantAct.__init__c Cs:|jj›d|j›d|j›d|j ¡d›d|j ¡d›d S)Nz(activation_bit=z, quant_mode: z , Act_min: z.2fz , Act_max: ú))r$r5r?rr;Úitemr<)r"r%r%r&Ú__repr__–sÿ ÿ þÿzQuantAct.__repr__c Cs¦|dur|n||}|jr†|jrJdƒ‚|jrJdƒ‚|j ¡}|j ¡} |  ¡ ¡dkr5| ¡ ¡dks9Jdƒ‚|j ¡dkrT|j  ¡dkrT|j||_|j | |_ n2|j dkrjt  |j|¡|_t  |j | ¡|_ n|j|j |d|j |_|j |j | d|j |_ |j s|dfS|dur”|jn|}|dur|j n|} t |j|| |jd |_|durº| ||j|j|j¡} n t |||j|j||¡} |j d¡} | | |jfS) Nz:percentile mode is not currently supported for activation.r>rz5NaN detected when computing min/max of the activationg¢&ú|”ç¾g¢&ú|”ç>éÿÿÿÿr )rA)ÚtrainingrBrAr(r*r,ÚisnanÚsumr;r<r@rrr-r?r=rCÚ FixedPointMulr Úview) r"r.Úpre_act_scaling_factorÚidentityÚidentity_scaling_factorZ specified_minZ specified_maxZx_actr;r<Z quant_act_intZcorrect_output_scaler%r%r&r3sH   "ÿ   ÿú zQuantAct.forward)r FNF)NNNNN©r5r6r7r8rrGr3r9r%r%r#r&r:rs ùr:cs:eZdZdZ d ‡fdd„ Z‡fdd „Zdd d „Z‡ZS)Ú QuantLineara8 Quantized version of `torch.nn.Linear`. Adds quantization-specific arguments on top of `torch.nn.Linear`. Args: weight_bit (`int`, *optional*, defaults to `8`): Bitwidth for the quantized weight. bias_bit (`int`, *optional*, defaults to `32`): Bitwidth for the quantized bias. per_channel (`bool`, *optional*, defaults to `False`): Whether or not to use channel-wise quantization. quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. Tré Fcs®tƒ ¡||_||_t t ||g¡¡|_|  dt  |j¡¡|  dt |j¡¡|r?t t |¡¡|_ |  dt  |j ¡¡||_ ||_ ||_||_||_ d|_tj|_dS)Nr Úfc_scaling_factorÚ bias_integerF)r rÚ in_featuresÚ out_featuresrrrrrrrÚbiasrrrAÚbias_bitrrr r!)r"rVrWrXrrYrArr#r%r&rës  zQuantLinear.__init__cs*tƒ ¡}d|›d|j›d|j›d}|S)Nú(z weight_bit=z , quant_mode=rE)r rGrr)r"Úsr#r%r&rGs zQuantLinear.__repr__Nc Cs |jstjj||j|jddfS|dur|jdksJdƒ‚|j}|j ¡}|j r=t j |ddd\}}t j |ddd\}}n|  ¡  d¡}|  ¡  d¡}t|j|||j ƒ|_| |j|j|j|j¡|_|j|}|jdurw| |j|jd|¡|_| dd¡}||} tjj| |j|jd||fS)N)rrX)r z«Input activation to the QuantLinear layer should be globally (non-channel-wise) quantized. Please add a QuantAct layer with `per_channel = True` before this QuantAct layerr )rÚoutFrH)rrr'ZlinearrrXÚshaper(r)rArr*r,r+r-rrTr!rr rYrUrM) r"r.Zprev_act_scaling_factorr/r0r1Ú_r2Zbias_scaling_factorÚx_intr%r%r&r3s0ÿ ÿ   þzQuantLinear.forward)TrrSFF©NrQr%r%r#r&rRÜs ÿ rRcs4eZdZdZd ‡fdd„ Zdd„Zd d d „Z‡ZS) ÚIntGELUa} Quantized version of `torch.nn.GELU`. Adds quantization-specific arguments on top of `torch.nn.GELU`. Args: quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. force_dequant (`str`, *optional*, defaults to `"none"`): Force dequantize the layer if either "gelu" or "nonlinear" is given. TÚnonecsjtƒ ¡||_|dvrt d¡d|_|jst ¡|_d|_d|_ gd¢|_ |j d|j d<dS) N)Ú nonlinearZgeluzForce dequantize geluFg à- ö?é)g]mÅþ²{Ò¿gçû©ñÒMü¿r ér) r rrÚloggerÚinforZGELUÚ activation_fnÚkÚconstÚcoeff)r"rÚ force_dequantr#r%r&r7s    zIntGELU.__init__cCsšt |jd|¡}t |jd|d¡}t |¡}t t |¡| ¡}|||d|}|d|jd}t |d|j¡}|d|j}||fS©Nr rer) rÚfloorrkÚsignr*ÚabsÚ floor_ster rj)r"r_Úscaling_factorÚb_intÚc_introZabs_intÚy_intr%r%r&Úint_erfGs zIntGELU.int_erfNcCs^|js | |¡dfS||}| |||j¡\}}d|}|||}||d}|||fS)Nçð?re)rrhrvri)r"r.rrr_Z sigmoid_intZsigmoid_scaling_factorZ shift_intr%r%r&r3Vs   zIntGELU.forward)Trbr`)r5r6r7r8rrvr3r9r%r%r#r&ra,s  racs:eZdZdZd ‡fdd„ Zdd„Zdd „Zd d „Z‡ZS) Ú IntSoftmaxaØ Quantized version of `torch.nn.Softmax`. Adds quantization-specific arguments on top of `torch.nn.Softmax`. Args: output_bit (`int`): Bitwidth for the layer output activation. quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. force_dequant (`str`, *optional*, defaults to `"none"`): Force dequantize the layer if either "softmax" or "nonlinear" is given. FrbcsŽtƒ ¡||_d|_||_|dvrt d¡d|_td|jd|_d|_ d|_ gd ¢|_ |j d |j d <|j d |j d <dS) NrS)rcÚsoftmaxzForce dequantize softmaxFé©rgvq à-æ¿é)gN„ª$ôëÖ?g¾Ã'|:ï?rwr rre) r rÚ output_bitÚmax_bitrrfrgr:ÚactÚx0rjÚcoef)r"r}rrlr#r%r&rrs   zIntSoftmax.__init__cCs~t ¡t |jd|¡}t |jd|d¡}Wdƒn1s%wY||||}|jd|d}||fSrm)rÚno_gradrnr)r"r_rrrsrtÚzr%r%r&Úint_polynomialƒs þzIntSoftmax.int_polynomialcCs¬t ¡t |j|¡}Wdƒn1swYt ||j|¡}t ||¡}|||}| ||¡\}}tj t |d|j|¡dd}|d|j}||fS)Nrer©r*) rr‚rnr€r,rjrqr r„Úclamp)r"r_rrZx0_intÚqÚrÚexp_intÚexp_scaling_factorr%r%r&Úint_exp‹s ÿ "zIntSoftmax.int_expc Cs¾|js tjj|dddfS||}|jddd\}}||}| ||¡\}}| ||¡\}}||}|jddd} t  d|j | ¡} t  || d|j |j ¡}dd|j }|||fS)NrH©rT)rÚkeepdimrer ) rrr'ryr,r‹rrKrqr r~r}) r"r.rrr_Z x_int_maxr^r‰rŠÚexpZ exp_int_sumÚfactorr%r%r&r3—s zIntSoftmax.forward)Frb) r5r6r7r8rr„r‹r3r9r%r%r#r&rxes   rxcs<eZdZdZd‡fdd„ Zdd„Zd d „Zdd d „Z‡ZS)Ú IntLayerNormaû Quantized version of `torch.nn.LayerNorm`. Adds quantization-specific arguments on top of `torch.nn.LayerNorm`. Args: output_bit (`int`, *optional*, defaults to `8`): Bitwidth for the layer output activation. quant_mode (`bool`, *optional*, defaults to `False`): Whether or not the layer is quantized. force_dequant (`str`, *optional*, defaults to `"none"`): Force dequantize the layer if either "layernorm" or "nonlinear" is given. rFrbcs’tƒ ¡||_||_t t |¡¡|_t t |¡¡|_ ||_ |dvr,t   d¡d|_ |  dt d¡¡||_d|_d|_t|j|j d|_dS)N)rcZ layernormzForce dequantize layernormFÚshiftr rSr{)r rÚnormalized_shapeÚepsrrrrrrXrrfrgrr}r~Údim_sqrtr:Z activation)r"r’r“r}rrlr#r%r&r¹s  zIntLayerNorm.__init__cCsžt ¡A|d}tj|ddd}t t |d|j¡¡ ¡ ¡}|j}t |j|¡|_t   dt |ƒ›dt |jƒ›¡WdƒdS1sHwYdS)NreT©ZaxisrzDynamic shift adjustment: z -> ) rr‚rKÚlog2Úsqrtr~Úceilr,r‘rfrgÚint)r"ruÚy_sq_intÚvar_intr‘Z shift_oldr%r%r&Ú set_shiftÌs """úzIntLayerNorm.set_shiftcCs:| |¡t |d|j¡}|d}tj|ddd}|S)z± This fallback function is called when overflow is detected during training time, and adjusts the `self.shift` to avoid overflow in the subsequent runs. reTr•)rœrqr r‘rrK)r"ruÚ y_int_shiftedršr›r%r%r&Úoverflow_fallbackÕs zIntLayerNorm.overflow_fallbackNcCs¬|js.|jddd}||}tj|dddd}|t |j|¡}||j|j}|dfS|jdurHtj|j dtj d}t |¡  |j ¡|_||}t  |jddd¡}||} t | d|j¡} | d} tj| ddd} |jr|  ¡d|jkr| | ¡} |  ¡d|jdksJdƒ‚t t | ¡¡d|j} t d| ¡}t | |d¡} |jd}|jj ¡|jj ¡}t ||¡}| |} ||j}| |}||fS) NreTr•)Zdtypegš™™™™™¹?zfError detected in overflow handling: `var_int` exceeds `self.max_bit` (the maximum possible bit width)li@)rÚmeanrr—r“rrXr”Útensorr]ÚfloatÚtoÚdeviceÚ round_ster rqr‘rKrIr,r~ržr(r))r"r.rrrŸÚyÚvarÚnr_Zmean_intrurršr›Zstd_intrrXZbias_intr%r%r&r3às@  ÿ  zIntLayerNorm.forward)rFrbr`) r5r6r7r8rrœržr3r9r%r%r#r&r¬s    rFc Cs€|jd}t|d|dƒ}t||dƒ}tj||dj}|dkr(|d}n tj| |dj }|s<| ¡}| ¡}||fS)aÆ Calculate the percentile max and min values in a given tensor Args: input (`torch.Tensor`): The target tensor to calculate percentile max and min. lower_percentile (`float`): If 0.1, means we return the value of the smallest 0.1% value in the tensor as percentile min. upper_percentile (`float`): If 99.9, means we return the value of the largest 0.1% value in the tensor as percentile max. output_tensor (`bool`, *optional*, defaults to `False`): If True, this function returns tensors, otherwise it returns values. Returns: `Tuple(torch.Tensor, torch.Tensor)`: Percentile min and max value of *input* rr g{®Gáz„?)ri)r]ÚroundrZkthvalueÚvaluesrF) ÚinputZlower_percentileZupper_percentileZ output_tensorZ input_lengthZ lower_indexZ upper_indexÚ upper_boundÚ lower_boundr%r%r&Úget_percentile_min_maxs  r­cCs¢t|jƒdkr| dddd¡}| dddd¡}nt|jƒdkr,| dd¡}| dd¡}n | d¡}| d¡}|rF| d|¡ |¡ ¡|St d|||¡S)a? Quantize single-precision input tensor to integers with the given scaling factor and zeropoint. Args: input (`torch.Tensor`): Single-precision input tensor to be quantized. scale (`torch.Tensor`): Scaling factor for quantization. zero_pint (`torch.Tensor`): Shift for quantization. inplace (`bool`, *optional*, defaults to `False`): Whether to compute inplace or not. Returns: `torch.Tensor`: Linearly quantized value of *input* according to *scale* and *zero_point*. érHr rerw)Úlenr]rMZmul_Zadd_Zround_rr¨)rªÚscaleÚ zero_pointÚinplacer%r%r&Úlinear_quantize5s   r³cCs²t ¡Kd|dd}|r-tjtj| ¡| ¡gdddd\}}tj|dd|}nt| ¡| ¡ƒ}tj|dd|}Wdƒ|SWdƒ|S1sRwY|S)a/ Compute the scaling factor with the given quantization range for symmetric quantization. Args: saturation_min (`torch.Tensor`): Lower bound for quantization range. saturation_max (`torch.Tensor`): Upper bound for quantization range. per_channel (`bool`, *optional*, defaults to `False`): Whether to or not use channel-wise quantization. Returns: `torch.Tensor`: Scaling factor that linearly quantizes the given range between *saturation_min* and *saturation_max*. rer rŒg:Œ0âŽyE>r…N)rr‚r,Ústackrpr†)Znum_bitsZsaturation_minZsaturation_maxrAr§r°r^r%r%r&r-Xs ( ÷ ú û õ r-c@ó(eZdZdZedd„ƒZedd„ƒZdS)rzw Class to quantize the given floating-point values using symmetric quantization with given range and bitwidth. cCsNtjd|jd}d|dd}t|||dd}t || |d¡}||_|S)a6 Args: x (`torch.Tensor`): Floating point tensor to be quantized. k (`int`): Quantization bitwidth. percentile_mode (`bool`): Whether or not to use percentile calibration. scale (`torch.Tensor`): Pre-calculated scaling factor for *x*. Note that the current implementation of SymmetricQuantFunction requires pre-calculated scaling factor. Returns: `torch.Tensor`: Symmetric-quantized value of *input*. g)r£rer F)r²)rr r£r³r†r°)Úctxr.rirr°r±r§Z new_quant_xr%r%r&r3}s zSymmetricQuantFunction.forwardcCsb|j}t|jƒdkr| dddd¡}nt|jƒdkr!| dd¡}n| d¡}| ¡|ddddfS)Nr®rHr re)r°r¯r]rMÚclone)r¶Ú grad_outputr°r%r%r&Úbackward—s zSymmetricQuantFunction.backwardN©r5r6r7r8Ú staticmethodr3r¹r%r%r%r&rxs  rc@rµ)rqz; Straight-through Estimator(STE) for torch.floor() cCó t |¡Sr`)rrn©r¶r.r%r%r&r3ªó zfloor_ste.forwardcCó| ¡Sr`©r·©r¶r¸r%r%r&r¹®ózfloor_ste.backwardNrºr%r%r%r&rq¥ó  rqc@rµ)r¤z; Straight-through Estimator(STE) for torch.round() cCr¼r`)rr¨r½r%r%r&r3¸r¾zround_ste.forwardcCr¿r`rÀrÁr%r%r&r¹¼rÂzround_ste.backwardNrºr%r%r%r&r¤³rÃr¤écCs®| ¡}| d¡}t | ¡ ¡¡\}}g}|D]}tt |d|¡j t d¡tj dƒ}|  |¡qt  |¡}t |ƒ|}t |¡ |j¡ |¡t |¡ |j¡ |¡fS)zü Decompose the scaling factor into mantissa and twos exponent. Args: scaling_factor (`torch.Tensor`): Target scaling factor to decompose. Returns: ``Tuple(torch.Tensor, torch.Tensor)`: mantisa and exponent rHreÚ1)Úrounding)ÚsizerMÚnpÚfrexpÚcpuÚnumpyr™ÚdecimalÚDecimalÚquantizeÚ ROUND_HALF_UPÚappendÚarrayr¡rZ from_numpyr¢r£)Zinputsr~Zshape_of_inputZoutput_mZoutput_eZtmp_mÚmZ int_m_shiftedr%r%r&Ú batch_frexpÁs "ÿ   þrÓc@s.eZdZdZe  ddd„ƒZedd„ƒZdS)rLaQ Function to perform fixed-point arithmetic that can match integer arithmetic on hardware. Args: pre_act (`torch.Tensor`): Input tensor. pre_act_scaling_factor (`torch.Tensor`): Scaling factor of the input tensor *pre_act*. bit_num (`int`): Quantization bitwidth. z_scaling_factor (`torch.Tensor`): Scaling factor of the output tensor. identity (`torch.Tensor`, *optional*): Identity tensor, if exists. identity_scaling_factor (`torch.Tensor`, *optional*): Scaling factor of the identity tensor *identity*, if exists. Returns: `torch.Tensor`: Output tensor(*pre_act* if *identity* is not given, otherwise the addition of *pre_act* and *identity*), whose scale is rescaled to *z_scaling_factor*. NcCs”t|jƒdkr dd„}ndd„}||_d|dd}t ¡¡||ƒ}|dur,||ƒ}||_t ||¡} | tj¡} | tj ¡ tj¡} | | } || ƒ} t | ƒ\} }|  tj¡|  tj¡}t |d|¡}|dur«t ||¡}| tj¡} | tj ¡ tj¡} | | } || ƒ} t | ƒ\}}| tj¡| tj¡}t |d|¡}||}t  | tj ¡| d|¡WdƒS1sÃwYdS)NrcSs|Sr`r%©r.r%r%r&Úsz'FixedPointMul.forward..cSs| ddd¡S)Nr rH)rMrÔr%r%r&rÕsrer r) r¯r]rOrr‚Úz_scaling_factorr¨ÚtypeÚdoubler¡rÓr†)r¶Zpre_actrNZbit_numrÖrOrPZreshaper§Zz_intZ_AZ_BZ new_scalerÒÚeÚoutputZwx_intÚm1Úe1Zoutput1r%r%r&r3ús<      $ßzFixedPointMul.forwardcCs8d}|jdur| ¡|j}| ¡|jdddd|dfSr`)rOr·rÖ)r¶r¸Z identity_gradr%r%r&r¹/s zFixedPointMul.backwardr4rºr%r%r%r&rLãs ù4rL)F)rÄ)rÌrËrÈrrZtorch.autogradrÚutilsrZ get_loggerr5rfÚModulerr:rRrarxrr­r³r-rrqr¤rÓrLr%r%r%r&Ús*    SjP9G e $ # - "