qi %ddlmZddlmZddlmZmZddlmZddl m Z ddl m Z erdd l Z erdd lmZej eZGd d e Zy ) ) annotations) TYPE_CHECKING)is_torch_availablelogging) SinqConfig) HfQuantizer)get_module_from_nameN)PreTrainedModelceZdZUdZdZded<dfd ZddZeddZ dZ dd Z dd Z dd Z dd Zd ZdZ d ddZ ddZxZS)SinqHfQuantizera HF v5 quantizer for SINQ. Modes: - method="sinq" (default): * weight-only SINQ * param-level ConversionOps (`SinqQuantize`) during load for pure language models (each Linear.weight is turned into a SINQLinear module) * module-level quantization after load for multimodal models - method="asinq": * A-SINQ (activation-aware) SINQ quantization Tbool requires_parameters_quantizationc Bt||fi|d|_d|_y)NF)super__init___normalized_device_str_do_param_level_sinq)selfquantization_configkwargs __class__s X/opt/pipecat/venv/lib/python3.12/site-packages/transformers/quantizers/quantizer_sinq.pyrzSinqHfQuantizer.__init__1s& ,7726#*/!cyNTrs ris_serializablezSinqHfQuantizer.is_serializable7srcyrrrs r is_trainablezSinqHfQuantizer.is_trainable:src|\tjjr!dtjji}nddi}tj d|d|S)Ncpuz:The device_map was not initialized. Setting device_map to zJ. If you want to use the model for inference, please set device_map='auto')torchcuda is_availablecurrent_deviceloggerinfo)r device_maps rupdate_device_mapz!SinqHfQuantizer.update_device_map>sb  zz&&( %**";";"=>  %[ KK))3 5[[  rc8|tj}||_|SN)r&bfloat16dtype)rr1s r update_dtypezSinqHfQuantizer.update_dtypeKs =NNE  rcddlm}|s tdtjj st jd|jd}t|tr?t|j}t|dkDrtdt|d|j j"d k(r|j$s t'd yy) Nr)is_sinq_availablezMThe 'sinq' package is not installed. Please install it with: pip install sinqzNo CUDA device is available. Quantization and inference will run on the CPU. Please note that this will significantly slow down inference speed and increase quantization time.r,r zkSinqHfQuantizer: multi-GPU device_map detected, but SINQ currently supports only a single CUDA device. Got z. Please use device_map=None.asinqzYou are using `method='asinq'` in the quantization config. Right now the calibrated version of SINQ is not supported in Hugging Face, please refer and use the official SINQ repository `to quantize a model with this method. )utilsr4 ImportErrorr&r'r(r*warningget isinstancedictsetvalueslen RuntimeErrorsortedrmethod pre_quantized ValueError)rargsrr4r,device_map_valuess rvalidate_environmentz$SinqHfQuantizer.validate_environmentQs- "mn nzz&&( NNB ZZ - j$ ' #J$5$5$7 8 $%)"##)*;#<"==Z\  # # * *g 5d>P>P: ?Q 5rc ddlm}|j}|t|j|j t|j ndddddt |j|S)zI Build the dict that SINQLinear expects as quant_config. r)sinq_base_quant_configNFr )nbits group_size quant_zero quant_scale view_as_floataxis tiling_moderA)sinq.sinqlinear_hfrHrAintrIrJstrrO)rcfgsinq_base_quant_config_fnrAs r_build_sinq_quant_dictz&SinqHfQuantizer._build_sinq_quant_dictms[ [(cii..1nn.Hs3>>*dCOO,  rc ddlm}|jry|jjdk(ry|j syt ||\}}|dk7ryt||}t|dd}|xr| } | S)a- Called per-parameter to decide whether to run `SinqQuantize` on it. - If `self.pre_quantized`, we do *not* quantize again (handled by SinqDeserialize instead). - For method="asinq": return False (ASINQ is not supported in Hugging Face). - For method="sinq": True only for SINQLinear.weight not in modules_to_not_convert. Note: After _process_model_before_weight_loading(), the modules are already SINQLinear, not nn.Linear. We check for SINQLinear modules that are not yet quantized (ready=False). r) SINQLinearFr5weightreadyT) rPrWrBrrArr r:getattr) rmodel param_namerrWmodule tensor_nameis_sinqis_readyresults rparam_needs_quantizationz(SinqHfQuantizer.param_needs_quantizations 2     # # * *g 5((25*E ( "VZ067D1)\ rcddlm}||S)z Return the ConversionOps used for param-level quantization (Sinq). The actual SINQLinear construction is in integrations/sinq.py. r) SinqQuantize)integrations.sinqrd)rrds rget_quantize_opsz SinqHfQuantizer.get_quantize_opss 5D!!rcbddlm}|jrddlm}|gddg||ggSgS)a4 If `pre_quantized=True`, interpret a checkpoint produced by SINQLinear.state_dict: .W_q .bias .meta via a WeightConverter + SinqDeserialize so that we reconstruct a SINQLinear module instead of a plain nn.Linear. r)WeightConverter)SinqDeserialize)z.W_qz.metaz.biasz.weight)source_patternstarget_patterns operations)core_model_loadingrhrBreri)rrhris rget_weight_conversionsz&SinqHfQuantizer.get_weight_conversionssE 9    ; % &/K / 56   rc bddlm}|j||jjxsg||_|jj dk(xr |j |_|j rdn|j|j}t|trEtt|jd}t|trd|}n.t|}n"t j"j%rdnd}|||j||j&||j }y) a Called on meta-initialized model, before loading any weights. For SINQ, we replace nn.Linear modules with empty SINQLinear modules here. The actual quantization happens later in SinqQuantize.convert() when weights are loaded. r)replace_with_sinq_linearsinqNrzcuda:zcuda:0r%)modules_to_not_convert quant_config compute_dtypedevicerB)rerpget_modules_to_not_convertrrrrArBrrUr:r;nextiterr=rQrRr&r'r(r1) rr[r,keep_in_fp32_modulesrrpsinq_quant_dict first_device device_strs r$_process_model_before_weight_loadingz4SinqHfQuantizer._process_model_before_weight_loadings A&*&E&E D,,CCIrL`' # %)$<$<$C$Cv$M$hVZVhVhRh!"&"4"4$$:U:UVZVnVn:o j$ 'Z%6%6%8 91=L,,$\N3  . %*ZZ%<%<%>EJ( #'#>#>(**,,  rc  ddlm}||S)aq Called after *all* weights have been loaded. For SINQ: 1. Move non-SINQLinear modules to GPU (embeddings, norms, lm_head, etc.) - SINQLinear modules already have GemLite buffers on GPU - We skip moving SINQLinear's W_q/meta to avoid memory duplication 2. Patch HF save/load methods for SINQ serialization r)patch_hf_pretrained_io) sinq.hf_ior)rr[rrs r#_process_model_after_weight_loadingz3SinqHfQuantizer._process_model_after_weight_loadings 6   r)rr)returnr)r1 torch.dtyperr)rNone)rSrrr;)r[r r\rRrrr/)r[r ryzlist[str] | None)r[r )__name__ __module__ __qualname____doc__r__annotations__rr propertyr"r-r2rFrUrbrfrnr}r __classcell__)rs@rrr!s .2$d10   8 $ D"B26 ) ) / ) Vrr) __future__rtypingrr6rrutils.quantization_configrbaser quantizers_utilsr r&modeling_utilsr get_loggerrr*rrrrrsJ# /220   H %dkdr