U 4A·fxuã@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é)Úloggingc s,eZdZdZd ‡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__©úK/tmp/pip-unpacked-wheel-zw5xktn0/transformers/models/ibert/quant_modules.pyr,s  zQuantEmbedding.__init__c Csº|js.tj ||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&ÚforwardMsBù ö  ÿù zQuantEmbedding.forward) NNrFFNrr F)NN)Ú__name__Ú __module__Ú __qualname__Ú__doc__rr3Ú __classcell__r%r%r#r&rsô!rcs4eZdZdZd ‡fdd„ Zdd„Zd 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|_ |jsŠ|  dt   d¡¡|  dt   d¡¡|  dt   d¡¡|j d8_ |jd7_ntdƒ‚dS)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?r@Z 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$r4r>rr:Úitemr;)r"r%r%r&Ú__repr__–s8ÿzQuantAct.__repr__c Cs°|dkr |n||}|jr|jr*tdƒ‚|jr8tdƒ‚|j ¡}|j ¡} |  ¡ ¡dkrl| ¡ ¡dksttdƒ‚|j  ¡dkrª|j  ¡dkrª|j ||_ |j | |_ nd|j dkrÖt  |j |¡|_ t  |j | ¡|_ n8|j |j |d|j |_ |j |j | d|j |_ |j s|dfS|dkr.|j n|}|dkrB|j n|} t|j|| |jd |_|dkr~| ||j|j|j¡} nt |||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 )r@)ÚtrainingrAÚAssertionErrorr@r(r*r,ÚisnanÚsumr:r;r?rrr-r>r<rBÚ 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&r3sT   ÿÿþ  ÿ ú zQuantAct.forward)r FNF)NNNNN©r4r5r6r7rrFr3r8r%r%r#r&r9rs ùr9cs8eZdZdZd ‡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Úbiasrrr@Úbias_bitrrr r!)r"rVrWrXrrYr@rr#r%r&rës  zQuantLinear.__init__cs*tƒ ¡}d|›d|j›d|j›d}|S)Nú(z weight_bit=z , quant_mode=rD)r rFrr)r"Úsr#r%r&rFs zQuantLinear.__repr__Nc Cs |js tjj||j|jddfS|dk r2|jdks:tdƒ‚|j}|j  ¡}|j rzt j |ddd\}}t j |ddd\}}n|  ¡ d¡}|  ¡ d¡}t|j|||j ƒ|_| |j|j|j|j¡|_|j|}|jdk rî| |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ÚoutFrG)rrr'ZlinearrrXÚshaperIr(r)r@rr*r,r+r-rrTr!rr rYrUrM) r"r.Zprev_act_scaling_factorr/r0r1Ú_r2Zbias_scaling_factorÚx_intr%r%r&r3s6ÿ ÿ   þ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Únonecsltƒ ¡||_|dkr(t d¡d|_|js8t ¡|_d|_d|_ dddg|_ |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 rdr) rÚfloorrjÚsignr*ÚabsÚ floor_ster ri)r"r_Úscaling_factorÚb_intÚc_intrnZabs_intÚy_intr%r%r&Úint_erfGs zIntGELU.int_erfNcCs^|js| |¡dfS||}| |||j¡\}}d|}|||}||d}|||fS)Nçð?rd)rrgrurh)r"r.rqr_Z sigmoid_intZsigmoid_scaling_factorZ shift_intr%r%r&r3Vs  zIntGELU.forward)Tra)N)r4r5r6r7rrur3r8r%r%r#r&r`,s r`cs: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. Fracstƒ ¡||_d|_||_|dkr4t d¡d|_td|jd|_d|_ d|_ d d d g|_ |j d |j d <|j d|j d <dS)NrS)rbÚsoftmaxzForce dequantize softmaxFé©rgvq à-æ¿égN„ª$ôëÖ?g¾Ã'|:ï?rvr rrd) r rÚ output_bitÚmax_bitrrerfr9ÚactÚx0riÚcoef)r"r|rrkr#r%r&rrs   zIntSoftmax.__init__c Csjt ¡2t |jd|¡}t |jd|d¡}W5QRX||||}|jd|d}||fSrl)rÚno_gradrmr€)r"r_rqrrrsÚzr%r%r&Úint_polynomialƒs  "zIntSoftmax.int_polynomialc Cs˜t ¡t |j|¡}W5QRXt ||j|¡}t ||¡}|||}| ||¡\}}tj t |d|j|¡dd}|d|j}||fS)Nrdr©r*) rrrmrr,rirpr rƒÚclamp)r"r_rqZx0_intÚqÚrÚexp_intÚexp_scaling_factorr%r%r&Úint_exp‹s  "zIntSoftmax.int_expc Cs¾|jstjj|dddfS||}|jddd\}}||}| ||¡\}}| ||¡\}}||}|jddd} t  d|j | ¡} t  || d|j |j ¡}dd|j }|||fS)NrG©rT)rÚkeepdimrdr ) rrr'rxr,rŠr~rKrpr r}r|) r"r.rqr_Z x_int_maxr^rˆr‰ÚexpZ exp_int_sumÚfactorr%r%r&r3—szIntSoftmax.forward)Fra) r4r5r6r7rrƒrŠr3r8r%r%r#r&rwes   rwcs<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. rFracs’tƒ ¡||_||_t t |¡¡|_t t |¡¡|_ ||_ |dkrXt   d¡d|_ |  dt d¡¡||_d|_d|_t|j|j d|_dS)N)rbZ layernormzForce dequantize layernormFÚshiftr rSrz)r rÚnormalized_shapeÚepsrrrrrrXrrerfrr|r}Údim_sqrtr9Z activation)r"r‘r’r|rrkr#r%r&r¹s  zIntLayerNorm.__init__c Csˆt ¡v|d}tj|ddd}t t |d|j¡¡ ¡ ¡}|j}t |j|¡|_t   dt |ƒ›dt |jƒ›¡W5QRXdS)NrdT©ZaxisrŒzDynamic shift adjustment: z -> ) rrrKÚlog2Úsqrtr}Úceilr,rrerfÚint)r"rtÚy_sq_intÚvar_intrZ 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. rdTr”)r›rpr rrrK)r"rtÚ y_int_shiftedr™ršr%r%r&Úoverflow_fallbackÕs  zIntLayerNorm.overflow_fallbackNcCs²|js\|jddd}||}tj|dddd}|t |j|¡}||j|j}|dfS|jdkrtj|j dtj d}t |¡  |j ¡|_||}t  |jddd¡}||} t | d|j¡} | d} tj| ddd} |jr&|  ¡d|jkr&| | ¡} |  ¡d|jdks&tdƒ‚t t | ¡¡d|j} t d| ¡}t | |d¡} |jd}|jj ¡|jj ¡}t ||¡}| |} ||j}| |}||fS) NrdTr”)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 rprrKrHr,r}rrIr(r))r"r.rqržÚyÚvarÚnr_Zmean_intrtrœr™ršZstd_intrŽrXZbias_intr%r%r&r3às@  ÿ  zIntLayerNorm.forward)rFra)N) r4r5r6r7rr›rr3r8r%r%r#r&r¬s    rFc Cs€|jd}t|d|dƒ}t||dƒ}tj||dj}|dkrP|d}ntj| |dj }|sx| ¡}| ¡}||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„?)rh)r]ÚroundrZkthvalueÚvaluesrE) ÚinputZlower_percentileZupper_percentileZ output_tensorZ input_lengthZ lower_indexZ upper_indexZ upper_boundZ lower_boundr%r%r&Úget_percentile_min_maxs  rªcCs¢t|jƒdkr0| dddd¡}| dddd¡}nr„)rrr,Ústackror…)Znum_bitsZsaturation_minZsaturation_maxr@r¦r­r^r%r%r&r-Xs (r-c@s(eZdZdZedd„ƒZedd„ƒZdS)rzw Class to quantize the given floating-point values using symmetric quantization with given range and bitwidth. cCsPt d¡ |j¡}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*. grdr F)r¯)rrŸr¡r¢r°r…r­)Úctxr.rhrr­r®r¦Z new_quant_xr%r%r&r3}s zSymmetricQuantFunction.forwardcCsb|j}t|jƒdkr&| dddd¡}n&t|jƒdkrB| dd¡}n | d¡}| ¡|ddddfS)Nr«rGr rd)r­r¬r]rMÚclone)r²Ú grad_outputr­r%r%r&Úbackward—s zSymmetricQuantFunction.backwardN©r4r5r6r7Ú staticmethodr3rµr%r%r%r&rxs  rc@s(eZdZdZedd„ƒZedd„ƒZdS)rpz; Straight-through Estimator(STE) for torch.floor() cCs t |¡S©N)rrm©r²r.r%r%r&r3ªszfloor_ste.forwardcCs| ¡Sr¸©r³©r²r´r%r%r&rµ®szfloor_ste.backwardNr¶r%r%r%r&rp¥s  rpc@s(eZdZdZedd„ƒZedd„ƒZdS)r£z; Straight-through Estimator(STE) for torch.round() cCs t |¡Sr¸)rr§r¹r%r%r&r3¸szround_ste.forwardcCs| ¡Sr¸rºr»r%r%r&rµ¼szround_ste.backwardNr¶r%r%r%r&r£³s  r£écCs®| ¡}| d¡}t | ¡ ¡¡\}}g}|D]6}tt |d|¡j t d¡tj dƒ}|  |¡q0t  |¡}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 rGrdÚ1)Úrounding)ÚsizerMÚnpÚfrexpÚcpuÚnumpyr˜ÚdecimalÚDecimalÚquantizeÚ ROUND_HALF_UPÚappendÚarrayr rZ from_numpyr¡r¢)Úinputsr}Zshape_of_inputZoutput_mZoutput_eZtmp_mÚmZ int_m_shiftedr%r%r&Ú batch_frexpÁs  "ÿ   þrÌc@s*eZdZdZeddd„ƒ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*. Nc CsŒt|jƒdkrdd„}ndd„}||_d|dd}t ¡B||ƒ}|dk rZ||ƒ}||_t ||¡} | tj¡} | tj ¡ tj¡} | | } || ƒ} t | ƒ\} }|  tj¡|  tj¡}t |d|¡}|dk rZt ||¡}| tj¡} | tj ¡ tj¡} | | } || ƒ} t | ƒ\}}| tj¡| tj¡}t |d|¡}||}t  | tj ¡| d|¡W5QR£SQRXdS)NrcSs|Sr¸r%©r.r%r%r&Úóz'FixedPointMul.forward..cSs| ddd¡S)Nr rG)rMrÍr%r%r&rÎrÏrdr 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}|jdk r| ¡|j}| ¡|jdddd|dfSr¸)rOr³rÐ)r²r´Z identity_gradr%r%r&rµ/s zFixedPointMul.backward)NNr¶r%r%r%r&rLãsù 4rL)F)F)F)r¼)rÄrÃrÀrrZtorch.autogradrÚutilsrZ get_loggerr4reÚModulerr9rRr`rwrrªr°r-rrpr£rÌrLr%r%r%r&Ús(    SjP9Ge $ # - "