o oiJ@sddlZddlmZmZmZmZmZmZddlZddl Z ddl m Z m Z m Z m Z ddlm mZddlZddlmZmZddlmZddlmZddlmZmZedd d ZGd d d e j jZGd dde j jZGddde j j Z!dedfddZ"Gddde j#Z$Gddde$Z%Gddde$Z&Gddde j j Z'ddZ(Gdd d e jZ)Gd!d"d"e jZ*Gd#d$d$e*Z+Gd%d&d&e*Z,Gd'd(d(e j#Z-Gd)d*d*e j#Z.Gd+d,d,e j#Z/dS)-N)AnyDictOptionalTypeVarUnionoverload)Tensordevicedtypenn) get_tile_inds undo_layout) QuantState)GlobalOptimManager)*INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING OutlierTracerTztorch.nn.Module)boundcseZdZdZ        ddededeedeed ed ed ed eed dffdd Z dddZ dddZ ded efddZ Z S)StableEmbeddinga Custom embedding layer designed to improve stability during training for NLP tasks by using 32-bit optimizer states. It is designed to reduce gradient variations that can result from quantization. This embedding layer is initialized with Xavier uniform initialization followed by layer normalization. Example: ``` # Initialize StableEmbedding layer with vocabulary size 1000, embedding dimension 300 embedding_layer = StableEmbedding(num_embeddings=1000, embedding_dim=300) # Reset embedding parameters embedding_layer.reset_parameters() # Perform a forward pass with input tensor input_tensor = torch.tensor([1, 2, 3]) output_embedding = embedding_layer(input_tensor) ``` Attributes: norm (`torch.nn.LayerNorm`): Layer normalization applied after the embedding. Methods: reset_parameters(): Reset embedding parameters using Xavier uniform initialization. forward(input: Tensor) -> Tensor: Forward pass through the stable embedding layer. N@Fnum_embeddings embedding_dim padding_idxmax_norm norm_typescale_grad_by_freqsparse_weightreturnc sJt||||||||| | tjj|| d|_t|dddidSa Args: num_embeddings (`int`): The number of unique embeddings (vocabulary size). embedding_dim (`int`): The dimensionality of the embedding. padding_idx (`Optional[int]`): Pads the output with zeros at the given index. max_norm (`Optional[float]`): Renormalizes embeddings to have a maximum L2 norm. norm_type (`float`, defaults to `2.0`): The p-norm to compute for the `max_norm` option. scale_grad_by_freq (`bool`, defaults to `False`): Scale gradient by frequency during backpropagation. sparse (`bool`, defaults to `False`): Computes dense gradients. Set to `True` to compute sparse gradients instead. _weight (`Optional[Tensor]`): Pretrained embeddings. )r weight optim_bits N) super__init__torchr LayerNormnormr get_instanceregister_module_override) selfrrrrrrrrr r  __class__K/home/ubuntu/.local/lib/python3.10/site-packages/bitsandbytes/nn/modules.pyr$3s  zStableEmbedding.__init__cCtjj|j|dSNr%r initxavier_uniform_r _fill_padding_idx_with_zeror*r-r-r.reset_parametersb z StableEmbedding.reset_parameterscCN|jdur%t|j|jdWddS1swYdSdSNrrr%no_gradr fill_r5r-r-r.r4m  "z+StableEmbedding._fill_padding_idx_with_zeroinputc CsDt||j|j|j|j|j|j}|t }| ||jj Sr0) F embeddingr rrrrrtor%get_default_dtyper'r r*r>embr-r-r.forwardrs zStableEmbedding.forward)NNrFFNNNrN)__name__ __module__ __qualname____doc__intrfloatboolrr$r6r4rE __classcell__r-r-r+r.rsB   / rcseZdZdZ       ddededeedeed ed ed ed eed ee ddffdd Z dddZ dddZ dedefddZ ZS) EmbeddingzS Embedding class to store and retrieve word embeddings from their indices. NrFrrrrrrrrr rc s8tj||||||||| d t|dddidSr)r#r$rr(r)) r*rrrrrrrrr r+r-r.r$s zEmbedding.__init__cCr/r0r1r5r-r-r.r6r7zEmbedding.reset_parameterscCr8r9r:r5r-r-r.r4r=z%Embedding._fill_padding_idx_with_zeror>c Cs&t||j|j|j|j|j|j}|Sr0)r?r@r rrrrrrCr-r-r.rEs zEmbedding.forward)NNrFFNNrF)rGrHrIrJrKrrLrMrr r$r6r4rErNr-r-r+r.rOsD    , rOcs|eZdZddddddejddf deejdeeded e d e d ej d ed de ddfddZ ddZ ddZddZddZe   d.dejdee efde d ed ddf ddZdd Zd/d!eeeee fd"e fd#d$Ze % % %d0d&ed!eeeefd'eee e fd"e def d(d)Zed1d&ed'ee e fd"e defd*d)Zed1d&ed+ed"e defd,d)Zfd-d)ZZS)2 Params4bitNF@Tfp4data quant_state blocksizecompress_statistics quant_type quant_storagemodule Linear4bit bnb_quantizedrc CsV|dur td}tj|||} || _|| _|| _|| _|| _| | _ || _ || _ | Sr9) r%emptyr_make_subclassrUrVrWrTrXr[rSrY) clsrS requires_gradrTrUrVrWrXrYr[r*r-r-r.__new__s zParams4bit.__new__cCs"|j}|j|d<|j|d<|S)NrSr_)__dict__copyrSr_r*stater-r-r. __getstate__s   zParams4bit.__getstate__cCs^|d|_|d|_|d|_|d|_|d|_|d|_|d|_|d|_|d |_dS) Nr_rUrVrWrTrSrXr[rY) r_rUrVrWrTrSrXr[rYrcr-r-r. __setstate__s        zParams4bit.__setstate__cCsHt|t|}|}||t|d|_t|d|_|S)NrTrS)typer`rerfrbdeepcopyrTrS)r*memo new_instancerdr-r-r. __deepcopy__s  zParams4bit.__deepcopy__cCs(t|t|}|}|||Sr0)rgr`rerf)r*rjrdr-r-r.__copy__s zParams4bit.__copy__cudaquantized_statsr_cKsttj|||}||_tj||d|_|jj|_|jj |_ |jj |_ d|_ |j |_||_|jdur8|j|j_|S)N)qs_dictr T)r%rr]rAr_r from_dictrTrUnestedrVrWr[r rXrY)r^rSrnr_r rYkwargsr*r-r-r.from_prequantized s      zParams4bit.from_prequantizedcCsZ|j|}tjj||j|j|j|j d\}}||_||_ |j dur(||j _ d|_ |S)N)rUrVrWrXT) rS contiguousrAbnb functional quantize_4bitrUrVrWrXrTrYr[)r*r ww_4bitrTr-r-r. _quantize&s  zParams4bit._quantizer non_blockingcCs |j|dur d|dS||dS)Nrm)r r{)rA)r*r r{r-r-r.rm6s zParams4bit.cuda.r*r cCdSr0r-r*r r r{r-r-r.rA9z Params4bit.tocCr|r0r-r*r r{r-r-r.rAAtensorcCr|r0r-r*rr{r-r-r.rADrc stjjj|i|\}}}}|dur|jdkr|js||S|jdur*|j|t t j|||d|j |j|j |j |j|jd}|S)Nrmr r r{)r_rTrUrVrWrX)r%_C_nn _parse_torgr[rzrTrArPr#r_rUrVrWrXr*argsrrr r r{convert_to_format new_paramr+r-r.rAGs    )FrmNNF....)rGrHrIr%uint8rrrrKrMstrr r`rerfrkrl classmethodrrrsrzrr rmrrrArNr-r-r+r.rPs      " &rPrY) Embedding4bitrZcCsrt|jdddur dSt|dddurtd|jjddks"Jt|jts2t|j|jdd|_|j|j_dS)NrTzhFP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first.T)rXr[) getattrr warningswarnshape isinstancerPrXrT)rYr-r-r.'fix_4bit_weight_quant_state_from_module]s rcsTeZdZdZddddejdffdd ZddZfd d Zd ej fd d Z Z S)rZa This class is the base module for the 4-bit quantization algorithm presented in [QLoRA](https://arxiv.org/abs/2305.14314). QLoRA 4-bit linear layers uses blockwise k-bit quantization under the hood, with the possibility of selecting various compute datatypes such as FP4 and NF4. In order to quantize a linear layer one should first load the original fp16 / bf16 weights into the Linear4bit module, then call `quantized_module.to("cuda")` to quantize the fp16 / bf16 weights. Example: ```python import torch import torch.nn as nn import bitsandbytes as bnb from bnb.nn import Linear4bit fp16_model = nn.Sequential( nn.Linear(64, 64), nn.Linear(64, 64) ) quantized_model = nn.Sequential( Linear4bit(64, 64), Linear4bit(64, 64) ) quantized_model.load_state_dict(fp16_model.state_dict()) quantized_model = quantized_model.to(0) # Quantization happens here ``` TNrRc sHt||||t|jjd||||d|_||_d|_d|_||_dS)aw Initialize Linear4bit class. Args: input_features (`str`): Number of input features of the linear layer. output_features (`str`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. Fr_rVrWrXrYN) r#r$rPr rS compute_dtypecompute_type_is_setrTrX) r*input_featuresoutput_featuresbiasrrVrWrXr r+r-r.r$s  zLinear4bit.__init__cCs|jtjtjfvr|j|_dS|jtjkrM|jtjkr0||jdkr0t dtj ddd|jtjkrO||jdkrQt dtj ddddSdSdSdS)NzInput type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference.ignorez .*inference.)messagezInput type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference or training speed.z.*inference or training) r r%float32bfloat16rfloat16numelrrrfilterwarnings)r*xr-r-r.set_compute_types   zLinear4bit.set_compute_typecsdt|||t|jdddur.|jjjddD]\}}|r#|n|||d|<qdSdS)zc save weight and bias, then fill state_dict with components of quant_state rTNT)packedzweight.)r#_save_to_state_dictrr rTas_dictitemsdetach)r* destinationprefix keep_varskvr+r-r.rs zLinear4bit._save_to_state_dictrcCst||jdur|jj|jkr|jj|j|j_|js%||d|_|j}|jdur3||j}|jdur:dn|j|j}tj ||j ||j j d|S)NT)rrT) rrr rSrArrrru matmul_4bitr trT)r*r inp_dtyperr-r-r.rEs   "zLinear4bit.forward) rGrHrIrJr%rr$rrrrErNr-r-r+r.rZns$%  rZc.eZdZdZdddejdffdd ZZS) LinearFP4z' Implements the FP4 data type. TNc t|||||d||dS)Q Args: input_features (`str`): Number of input features of the linear layer. output_features (`str`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. rRNr#r$r*rrrrrVrXr r+r-r.r$zLinearFP4.__init__rGrHrIrJr%rr$rNr-r-r+r.rsrcr) LinearNF4a"Implements the NF4 data type. Constructs a quantization data type where each bin has equal area under a standard normal distribution N(0, 1) that is normalized into the range [-1, 1]. For more information read the paper: QLoRA: Efficient Finetuning of Quantized LLMs (https://arxiv.org/abs/2305.14314) Implementation of the NF4 data type in bitsandbytes can be found in the `create_normal_map` function in the `functional.py` file: https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L236. TNc r)rnf4Nrrr+r-r.r$rzLinearNF4.__init__rr-r-r+r.r src seZdZ     ddeejdeejdeejfddZfd d Zd d Ze dde dee e e fdee eefdede f ddZe dde de eefdede fddZe dde dedede fddZfddZZS) Int8ParamsNTFrSCBSCBcCs8|dur td}tj|||}||_||_||_|Sr9)r%r\rr]rrhas_fp16_weights)r^rSr_rrrobjr-r-r.r`7s zInt8Params.__new__csN|jr t|S|j|}tj|\}}}||_||_ ||_ |Sr0) rr#rmrSrthalfrurvint8_vectorwise_quantrr)r*r Brr_r+r-r.rmGs zInt8Params.cudac CsDt|jt|t|j||j|jt|j|t|j|d}|S)N)rSr_rrr) rgr`rbrhrSr_rrr)r*rirjr-r-r.rkTs   zInt8Params.__deepcopy__.r*r r r{rcCr|r0r-r}r-r-r.rA`r~z Int8Params.tocCr|r0r-rr-r-r.rAhrrcCr|r0r-rr-r-r.rAkrcsztjjj|i|\}}}}|dur#|jdkr#|jjjdkr#||Stt j |||d|j |j d}|j |_ |j|_|S)Nrmcpur)r_r)r%rrrrgrSr rmrr#rAr_rrrrr+r-r.rAns  )NTFNNrr)rGrHrIrr%rr`rmrkrrrrKr r rrMrArNr-r-r+r.r6sF    &rc Cs||d}|durdS||dd}t|tjr!|}t|tr1|tvr1td|t|tr>|tvr>t|}|dkrTt ||j } t || ||d<dSdS)Nr weight_formatrowz'Expected supported weight format - got ) getpoprr%ritemrKr ValueErrorr r r ) state_dictrlocal_metadatastrict missing_keysunexpected_keys error_msgsr r tile_indicesr-r-r.maybe_rearrange_weights  rcs<eZdZdZd fdd ZddZdedefd d ZZS) Embedding8bita This class implements [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm for embedding layer Quantization API is similar to Linear8bitLt: ```python import torch import torch.nn as nn from bitsandbytes.nn import Embedding8bit fp16_module = nn.Embedding(128, 64) int8_module = Embedding8bit(128, 64) int8_module.load_state_dict(fp16_module.state_dict()) int8_module = int8_module.to(0) # Quantization happens here ``` Ncs8tj||||d|jjj|_t|jjddd|_dS)Nr r Frr_)r#r$r rSr r)r*rrr r r+r-r.r$s zEmbedding8bit.__init__cCtd)Nz.Saving Embedding8bit module is not implementedNotImplementedErrorr*rrrr-r-r.rz!Embedding8bit._save_to_state_dictr>rcCst|jds td|jj}|jj}|j|j|jfksJ|j|jfks&Jt ||}t || |jd}||d}| |j S)NrzKEmbedding layer is not quantized. Please call .cuda() or .to(device) first.rg_@) hasattrr RuntimeErrorrSrrrrr?r@viewrAr )r*r>rows row_statscompressed_outputcompressed_output_statsoutputr-r-r.rEs    zEmbedding8bit.forward)NN) rGrHrIrJr$rrrErNr-r-r+r.rs rcsTeZdZdZddejdffdd ZdefddZd d Z ded efd d Z Z S)ra3 This is the base class similar to Linear4bit. It implements the 4-bit quantization algorithm presented in [QLoRA](https://arxiv.org/abs/2305.14314) for embeddings. Quantization API is similar to Linear4bit: ```python import torch import torch.nn as nn from bitsandbytes.nn import Embedding4bit fp16_module = nn.Embedding(128, 64) quantized_module = Embedding4bit(128, 64) quantized_module.load_state_dict(fp16_module.state_dict()) quantized_module = quantized_module.to(0) # Quantization happens here ``` NrRcsntj||||d|jjj|_t|jjdd|||d|_|jj}||dkr5td|d|ddSdS)NrFrrzEmbedding size z is not divisible by block size z#. This will lead to slow inference.) r#r$r rSr rPrUrr)r*rrr rWrXr rUr+r-r.r$s   zEmbedding4bit.__init__r>c CsR|j|jjjdks J|jjtj|j|jdd}tj j j ||j|jd|ddd}|j | |jddfksCJ|j|jj}|jjj}|j |j|fksZJtj j j ||j||dd}|j | |fkswJt|jj}||_tg|j |jR|_ tj ||}|j g|j |jRksJ||jS)Nrrr r>r)rr rTrUrSrr%rrr rvr@rrabsmaxrbrhSizerudequantize_4bitrAr ) r*r> w_4bit_uint8 output_4bitblocks_per_embr output_absmaxoutput_quant_staterr-r-r. _forward_with_partial_dequantizes6$    z.Embedding4bit._forward_with_partial_dequantizecCr)Nz.Saving Embedding4bit module is not implementedrrr-r-r.rrz!Embedding4bit._save_to_state_dictrcCsVt||j|jjjdkr||Stj|jj |jj}t j jj ||d |jS)Nrr)rrr rTrUrrurvrrSr%r r@rAr )r*r>dequantized_weightr-r-r.rEs zEmbedding4bit.forward) rGrHrIrJr%rr$rrrrErNr-r-r+r.rs!rc&eZdZdejdffdd ZZS) EmbeddingFP4Nctj|||d||ddS)NrRr rWrXr rr*rrr rXr r+r-r.r$) zEmbeddingFP4.__init__rGrHrIr%rr$rNr-r-r+r.r( rcr) EmbeddingNF4Ncr)Nrrrrr+r-r.r$<rzEmbeddingNF4.__init__rr-r-r+r.r;rrcsfeZdZdZ     ddedeffdd Zfd d Zfd d Zd dZde j fddZ Z S) Linear8bitLtaZ This class is the base module for the [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm. To read more about it, have a look at the paper. In order to quantize a linear layer one should first load the original fp16 / bf16 weights into the Linear8bitLt module, then call `int8_module.to("cuda")` to quantize the fp16 weights. Example: ```python import torch import torch.nn as nn import bitsandbytes as bnb from bnb.nn import Linear8bitLt fp16_model = nn.Sequential( nn.Linear(64, 64), nn.Linear(64, 64) ) int8_model = nn.Sequential( Linear8bitLt(64, 64, has_fp16_weights=False), Linear8bitLt(64, 64, has_fp16_weights=False) ) int8_model.load_state_dict(fp16_model.state_dict()) int8_model = int8_model.to(0) # Quantization happens here ``` TNrrcsht||||t|_||_||j_||j_|dkr#|s#d|j_t |j j ||d|_ | t dS)ay Initialize Linear8bitLt class. Args: input_features (`int`): Number of input features of the linear layer. output_features (`int`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. rTrN)r#r$ru MatmulLtStaterdindex thresholdruse_poolrr rS"_register_load_state_dict_pre_hookr)r*rrrrrrr r+r-r.r$ns  zLinear8bitLt.__init__c st|||d}t|j|}t|j|}||}|d}|jjsW|dur=|r+|n|||<tjdtj d||<dS|durY|rE|n|||<tjdtj d||<dSdSdS)Nrrr)r ) r#rrr rdrrr%rr) r*rrrscb_nameparam_from_weightparam_from_statekey_name format_namer+r-r.rs   z Linear8bitLt._save_to_state_dictc st|||||||t|}|D]4} | t|d} | dkrF|jjdur*td|| } |jj| |jjdurA|jj|j_| | qdS)NrzLoading a quantized checkpoint into non-quantized Linear8bitLt is not supported. Please call module.cuda() before module.load_state_dict()) r#_load_from_state_dictlistlenr rrcopy_rdremove) r*rrrrrrrunexpected_copykey input_name input_paramr+r-r.r s0      z"Linear8bitLt._load_from_state_dictcC,|jj|j_|jj|j_d|j_d|j_dSr0r rrdrr5r-r-r.init_8bit_state   zLinear8bitLt.init_8bit_statercCs|j|j_|jjdur||jdur%|jj|jkr%|jj |j|j_t j ||j|j|jd}|jj sA|jjdurA|jj|j_|SN)rrd) trainingrd is_trainingr rrrr rSrArumatmulrr*routr-r-r.rEs   zLinear8bitLt.forward)TTrNN) rGrHrIrJrKr$rr rr%rrErNr-r-r+r.rNs # " 'rcs6eZdZd fdd ZddZddZd d ZZS) OutlierAwareLinearTNcs"t||||d|_d|_dSr)r#r$ outlier_dim is_quantized)r*rrrr r+r-r.r$s zOutlierAwareLinear.__init__cCr)NzJPlease override the `forward_with_outliers(self, x, outlier_idx)` functionr)r*r outlier_idxr-r-r.forward_with_outliersrz(OutlierAwareLinear.forward_with_outlierscCr)NzEPlease override the `quantize_weights(self, w, outlier_idx)` functionr)r*rxrr-r-r.quantize_weightrz"OutlierAwareLinear.quantize_weightcCsf|jdurt}|std||j}||_|js1||j|j}|jj |d|_dSdS)NzTPlease use OutlierTracer.initialize(model) before using the OutlierAwareLinear layerT) rrr(is_initializedprint get_outliersr rr!rSr )r*rtracerrrxr-r-r.rEs   zOutlierAwareLinear.forward)TN)rGrHrIr$r r!rErNr-r-r+r.rs rcs:eZdZ      d fdd ZddZd d ZZS) SwitchBackLinearBnbTFrNc sft||||t|_||_||j_||j_||j_|dkr'|s'd|j_ t |j j ||d|_ dS)NrTr) r#r$rurrdrrrmemory_efficient_backwardrrr rS) r*rrrrr'rrr r+r-r.r$s  zSwitchBackLinearBnb.__init__cCrr0rr5r-r-r.rrz#SwitchBackLinearBnb.init_8bit_statecCsF|j|j_|jjdur|tj||jd|jd|j }dSr) rrdrr rrru matmul_mixedrrrr-r-r.rEs  (zSwitchBackLinearBnb.forward)TTFrNN)rGrHrIr$rrErNr-r-r+r.r&sr&)0rbtypingrrrrrrrr%rr r r torch.nn.functionalrvr? bitsandbytesru bitsandbytes.autograd._functionsr r bitsandbytes.functionalrbitsandbytes.optimrbitsandbytes.utilsrrrrOr ParameterrPrLinearrZrrrrrrrrrrr&r-r-r-r.s<    jO y$+I/d