biddlZddlZddlZddlZdZGddZddZdedeeeffdZ dd Z d Z d Z dd d ddZ e jDcic]\}}|| c}}Zdej fdZycc}}w)Nc rt|tjjsJtj }|j |j}||jvrt|j}|jj||jj|t|jdkDrv|jdjdkDr9|jdj|jj dksJ|jd|j|<y|dj#d|dj d}t|dd}tj$|ddkDj't)t+t|dj dz }tj,|dkDd}tj.||gj1}||j|<y|j2D]}|j5y)Nr) reduction_dimzscore)dim) isinstancetorchnnLinear OutlierTracer get_instance get_hvalueweighthvalue2outlier_idxfind_outlier_dimsoutliersappendhvalueslennumelmaxshapeviewabssumlistrangewherecatuniquehooksremove) moduleinputtracerhvalue outlier_idxmergeddims outlier_idx2hooks M/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/bitsandbytes/utils.py outlier_hookr0s fehhoo .. .  ' ' )F   v}} -F V...' 6 {+f% v ! #r"((*Q.r*..06==3F3Fq3IIII060CF % %f -1X]]2uQx~~b'9:F+F!ANKIIeAh'!+00T%E!HNN@SVW@W:X5Y0ZD ;;tax03L))[,$?@GGIK0;F % %f -LL D KKM c>eZdZdZdZdZdZdZdZe dZ y)rNctd)NzCall get_instance() instead) RuntimeErrorselfs r/__init__zOutlierTracer.__init__,s899r1c<d|_d|_g|_g|_i|_d|_g|_|jD]X\}}t|tjjs+|j j|jtZy)NT)last_wcurrent_outlier_dimsrrr initializedr$ named_modulesr r r rrregister_forward_pre_hookr0)r6modelnms r/ initializezOutlierTracer.initialize/s $(!  "$ '') MDAq!UXX__- !!!"="=l"KL Mr1ct|ddS)Nr;F)getattrr5s r/is_initializedzOutlierTracer.is_initialized<st]E22r1cR|jjjSN)datastoragedata_ptr)r6rs r/rzOutlierTracer.get_hvalue?s{{""$--//r1c|js tdy|j|}||jvr|j|Sy)Nz$Outlier tracer is not initialized...)rDprintrr)r6rr)s r/ get_outlierszOutlierTracer.get_outliersBsJ""$ 8 9( T,, ,**62 2r1c^|j|j||_|jSrF) _instance__new__)clss r/rzOutlierTracer.get_instanceLs& == KK,CM}}r1) __name__ __module__ __qualname__rNr7rArDrrL classmethodrr1r/rr)s5I: M30r1rc|r?tjd|jd|f|jj S|j |}|j }|j }||z |z }|+tj|j|d\} } | Stj||kDd} | S)Nrr)sizedevice)kr ) r randintrrXlongstdmeantopkrr!) rrrr^rdmr\stdmstdstdzstdvalidxs r/rrSs }}Q QtgfmmTYY[[ **] #C 88:D WWYF $J& D ::cggi4Q7S Jkk$-(+ Jr1command_stringreturnc2dfd}||\}}||fS)Nc&td|DS)Nc3ZK|]#}|jdj%yw)zUTF-8N)decodestrip).0 to_decodes r/ z6execute_and_return.._decode..gs$a9Y%%g.446as)+)tuple)subprocess_err_out_tuples r/_decodez#execute_and_return.._decodefsaH`aaar1ctjtj|tjtjj S)N)stdoutstderr) subprocessPopenshlexsplitPIPE communicate)rerqs r/&execute_and_return_decoded_std_streamszBexecute_and_return..execute_and_return_decoded_std_streamsisA    N+!! km   r1rU)rer{std_outstd_errrqs @r/execute_and_returnr~es*b >nMGW G r1c4|jD]\}}tt|jdkDrt |||||t |t jjs`||vse|j|}||j|j|jdu|j|<|r<|j|j|_ |j|j|_ |t||d}||||S)a Replace linear modules with a new Linear module. Parameters: model (`torch.nn.Module`): Input model or `torch.nn.Module` as the function is run recursively. linear_replacement (`torch.nn.Module`): The linear module that replaces the old one. Only expects standard arguments. If other arguments need to be passed, use a lambda. skip_modules (`List[str]`, *optional*, defaults to `lm_head`): List of modules names not to convert. Defaults to `lm_head`. copy_weights (`bool`): Copy the weights from the old linear module to the new one post_processing_function (`str`): A function name of the replacement linear class that is called after processing. rN)named_childrenrrchildrenreplace_linearr r r r_modules in_features out_featuresbiasrrC) r>linear_replacement skip_modules copy_weightspost_processing_functionnamer& old_modulefuncs r/rrvs.,,.! f tFOO%& '! + 6#5|\Sk l fehhoo .4|3K-J#5""## 4'$ENN4 .8.?.?t$+,6OOt$)'3v'?F#L%!& Lr1ctj|}|jd}tjt |tj }|S)z Pack a dictionary into a torch tensor for storing quant_state items in state_dict. Parameters: - source_dict: The dictionary to be packed. Returns: A torch tensor containing the packed data. utf-8)dtype)jsondumpsencoder tensorruint8) source_dictjson_str json_bytes tensor_datas r/pack_dict_to_tensorrs?zz+&H)J,,tJ/u{{CK r1ct|jj}|jd}t j |}|S)z Unpack a torch tensor into a Python dictionary. Parameters: - tensor_data: The torch tensor containing the packed data. Returns: A Python dictionary containing the unpacked data. r)bytescpunumpyrjrloads)rrr unpacked_dicts r/unpack_tensor_to_dictrsB{(..01J  )HJJx(M r1rr)rowcol32 col_turing col_amperetc|jjdk(rtjj y|jjdk(rtj j yy)Ncudaxpu)rXtyper r synchronizer)rs r/sync_gpursGxx}}  %   r1)rg@NF))lm_headFN)rrwrur r0rrstrror~rrr"LINEAR_8BIT_WEIGHTS_FORMAT_MAPPINGitems*INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPINGTensorr)rrcs00r/rs  B''T$suS#X(! *Z""./![\%]"KmKsKsKu-vKT3c4i-v*   .ws A8