a h@sdZddlZddlmZddlmZddlmZmZddl Z ddl Z ddl m Z ddl m Z dd lmZdd lmZdd lmZmZmZmZmZmZd d lmZeeZdaddZGddde j j!Z"d,ddZ#d-ddZ$Gddde j%Z&Gddde j%Z'GdddeZ(eGdddeZ)eedd Gd!d"d"eZ*eed#d Gd$d%d%eZ+eGd&d'd'e)Z,ed(d Gd)d*d*e)e Z-gd+Z.dS).zPyTorch RWKV model.N) dataclass)Path)OptionalUnion)nn)GenerationMixin)GradientCheckpointingLayer)PreTrainedModel) ModelOutputauto_docstringis_bitsandbytes_availableis_ninja_availableis_torch_cuda_availablelogging) RwkvConfigcsddlm}ttjjjddfdddD}tdurNtj|krNdSt d|d d d d d ddd|g}|d||t t j k|da|t_dS)Nr)loadZkernelsrwkvcsg|] }|qSr).0fZ kernel_folderrb/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/rwkv/modeling_rwkv.py 5z(load_wkv_cuda_kernel..)z wkv_op.cppz wkv_cuda.cuzwkv_cuda_bf16.cuz2Loading CUDA kernel for RWKV at context length of .z -res-usagez--maxrregcount 60z--use_fast_mathz-O3z -Xptxas -O3z--extra-device-vectorizationz-DTmax=Zwkv_)namesourcesverboseZextra_cuda_cflags) Ztorch.utils.cpp_extensionrr__file__resolveparentrwkv_cuda_kernelmax_seq_lengthloggerinforZ get_verbosityDEBUG)context_lengthZ load_kernelZcuda_kernel_filesflagsrrrload_wkv_cuda_kernel/s*   r*c@s(eZdZedddZedddZdS) RwkvLinearAttentionNFc Cs|\}}} |tjkr0td|dtjd|| t| ddkrhtd|d| dt| dd |j|_|jjd ks|jjd ks|jjd ks|jjd krtd t |  }|jt j kr| }| }| }| }| }| }t j|t jd } |s|dur|dur^t j|| d t j|jt jd}|dddddfd8<nt jdd|Ddd }|jt jkrtj} ntj} | ||||| |n*|jt jkrtjntj} | ||||| |||||| |dur ddt j|d ddD}| |j|fS)NzCannot process a batch with z+ tokens at the same time, use a maximum of z with this model. rzThe product of batch size (z) and hidden size (z") needs to be a round multiple of rcudazUCalling the CUDA kernel for wkv attention requires all tensors to be on CUDA devices. memory_formatr)dtypedevicer/籡*GcSsg|]}|dqSr2Z unsqueezersrrrr~rz/RwkvLinearAttention.forward..)dimcSsg|]}|dqSr4)Zsqueezer6rrrrr)sizer#r$ ValueErrorminr0 input_dtyper1typetorchexpfloat contiguousfloat16 empty_likecontiguous_formatzerosfloat32catbfloat16Zforward_with_state_bf16Zforward_with_stateZ forward_bf16forwardZsave_for_backwardchunkto) ctx time_decay time_firstkeyvaluestate return_stateZ batch_sizeZseq_len hidden_sizeoutputZ forward_funcrrrrIPsl         zRwkvLinearAttention.forwardc Cs|j}|j\}}}}}tj|tj|tjkr0tjntjd} tj|tjd} tj|tjd} tj|tjd} |tjkr||}|tjkrt j nt j } | |||||| | | | | | || || || |ddfS)N)r/r0r.)r<Z saved_tensorsr>rCrDrHrFrBr@r#Z backward_bf16backwardrArK)rLZg_outputZg_stater<rMrNrOrPrTZ g_time_decayZ g_time_firstZg_keyZg_valueZ backward_funcrrrrUs@ zRwkvLinearAttention.backward)NF)N)__name__ __module__ __qualname__ staticmethodrIrUrrrrr+Os >r+FcCs|\}}}t|}|durztj|dddftjd} tj|dddftjd} tj|dddftjdd} n |\} } } t| }t|D]} |dd| f} |dd| f}t| | |}t| |}t| ||}|| ||}|| |}|||j |dd| f<t| || }t| ||}t| |}|| ||} || |} |} q|s|dur| | | g}||fS)Nr)r0r3) r9r>Z zeros_likerFr?ranger@maximumrKr0)rMrNrOrPrQrR_Z seq_lengthrTZ num_stateZ den_stateZ max_stateZ current_index current_key current_valueZmax_for_outpute1e2 numerator denominatorZ max_for_staterrrrwkv_linear_attention_cpus4 "      rccCsdtdd||||fD}|ddk}tdus8|s8|rLt||||||dSt||||||SdS)Ncss|]}|jjdkVqdS)r-N)r1r=)rtrrr rz(rwkv_linear_attention..rrQrR)anyr9r#rcr+apply)rMrNrOrPrQrRZno_cudaZ one_tokenrrrrwkv_linear_attentions rics2eZdZd fdd Zd ddZd dd ZZS) RwkvSelfAttentionrcsDt||_tduo"tj|jk}tr`tr`|s`zt|jWnt y^t dYn0||_ |j }|jdur||jn|}||_tt||_tt||_ttdd||_ttdd||_ttdd||_td|_tj||dd|_tj||dd|_tj||dd|_tj||dd|_dS)Nz9Could not load the custom CUDA kernel for RWKV attention.rrrrFbias)super__init__configr#r$r(rrr* Exceptionr%r&layer_idrSattention_hidden_sizer Parameterr>emptyrMrN time_mix_keytime_mix_valuetime_mix_receptance ZeroPad2d time_shiftLinearrOrP receptancerT)selfrqrsZ kernel_loadedrSrt __class__rrrps.   zRwkvSelfAttention.__init__NcCs|ddkr4|dur4|ddddd|jf}n:||}|durn|ddddd|jf|dddf<||j|d|j}||j|d|j}||j|d|j}||}||}t | |}|dur |dddf|ddddd|jf<||||fSNrrrl) r9rsr{rwrxryrOrPr>sigmoidr})r~hiddenrQshiftedrOrPr}rrrextract_key_values (   (z#RwkvSelfAttention.extract_key_valueFc sj||d\}}}}|dur.r2rfrrr)rtuplerirMrNrsrT) r~rrQ use_cacher}rOrPZ layer_staterrrrrI"s*    zRwkvSelfAttention.forward)r)N)NF)rVrWrXrprrI __classcell__rrrrrjs rjcs(eZdZdfdd ZdddZZS) RwkvFeedForwardrcst||_||_|j}|jdur,|jnd|j}td|_t t dd||_ t t dd||_ tj||dd|_tj||dd|_tj||dd|_dS)NrrkrFrm)rorprqrsrSintermediate_sizerrzr{rur>rvrwryr|rOr}rP)r~rqrsrSrrrrrp7s  zRwkvFeedForward.__init__NcCs |ddkr4|dur4|ddddd|jf}n:||}|durn|ddddd|jf|dddf<||j|d|j}||j|d|j}tt||}| |}t | |}|dur|dddf|ddddd|jf<|||fSr) r9rsr{rwryr>ZsquareZrelurOrPrr})r~rrQrrOr}rPrrrrIHs ( (zRwkvFeedForward.forward)r)NrVrWrXrprIrrrrrr6srcs&eZdZfddZdddZZS) RwkvBlockcsvt||_||_|dkr2tj|j|jd|_tj|j|jd|_ tj|j|jd|_ t |||_ t |||_dS)Nr)eps)rorprqrsr LayerNormrSZlayer_norm_epsilonpre_lnln1ln2rj attentionr feed_forward)r~rqrsrrrrp]s  zRwkvBlock.__init__NFcCsz|jdkr||}|j||||d\}}||}|j|||d\}}||}||f}|rn||f7}n|d7}|S)Nr)rQrrr)rsrrrrr)r~rrQroutput_attentionsrroutputsrrrrIks   zRwkvBlock.forward)NFFrrrrrr\s rc@s@eZdZUeed<dZdgZddgZdZdZ e j ddd Z d S) RwkvPreTrainedModelrqrrrMrNT)modulec st|tr<|j}|jj}|jj|j||dd||}tjfddt D|j j |j j d}|ddddf}fddt D}tj||j j |j j d}tjddt D|jj |jj dd }||j _t|jtd ||j_t|||j _t||d |j_t|d ||j_nt|tr|j}|jj}|jjd||}tjfd dt D|j j |j j d}|ddddf}t|||j _t|||j_nt|tjrn|jjj}d} d} |jdur|jj|d |dkr*t|d |d} |d |jjkrR|d|jjkrRd } | | 9} tjj |j| d npt|tj!r|jjj}dtt"|d |d} tjj |j| d n(t|tj#r|jj$d|jjdS)zInitialize the weights.rg?csg|] }|qSrrrirSrrrrz5RwkvPreTrainedModel._init_weights..r0r1Ncs,g|]$}dd|dddqS)rgffffff?g?r)rh)rt ratio_0_to_1rrrscSsg|]}|dddqS)rrrrrrrrrg?g333333?csg|] }|qSrrrrrrrrr)gaing-C6?)% isinstancerjrsrqnum_hidden_layersrSrtr>ZtensorrZrwr0r1rMrNdataZ ones_likemathlogpowrxryrrr|weightshapernZzero_sqrt vocab_sizeinitZ orthogonal_ EmbeddingmaxrZfill_) r~rrsrZratio_1_to_almost0Z time_weightZ decay_speedZzigzagrrscaler)rtrSrr _init_weightss|          $ z!RwkvPreTrainedModel._init_weightsN) rVrWrXr__annotations__Zbase_model_prefixZ_no_split_modulesZ_keep_in_fp32_modulesZsupports_gradient_checkpointingZ _is_statefulrModulerrrrrr~s rz+ Class for the RWKV model outputs. )Z custom_introc@sneZdZUdZdZeejed<dZ ee ejed<dZ ee ejdfed<dZ ee ejdfed<dS) RwkvOutputa  state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. Nlast_hidden_staterQ. hidden_states attentions)rVrWrX__doc__rrr> FloatTensorrrQlistrrrrrrrrs rzK Base class for causal language model (or autoregressive) outputs. c@seZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejdfed<dZee ejdfed<dS) RwkvCausalLMOutputap loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. NlosslogitsrQ.rr)rVrWrXrrrr>rrrrQrrrrrrrrrs  rc seZdZfddZddZddZedeej eej eej ee ej ee ee ee ee e eefd d d Zd d Zd dZZS) RwkvModelcsdttjj|_tfddtj D|_ t j|_ d|_ d|_|dS)Ncsg|]}t|dqS)r)r)ridxrqrrrrz&RwkvModel.__init__..F)rorprrrrS embeddingsZ ModuleListrZrblocksrln_outlayers_are_rescaledgradient_checkpointing post_initr~rqrrrrps  zRwkvModel.__init__cCs|jSrrrrrrget_input_embeddingsszRwkvModel.get_input_embeddingscCs ||_dSrrr~Znew_embeddingsrrrset_input_embeddingsszRwkvModel.set_input_embeddingsN) input_idsattention_mask inputs_embedsrQrroutput_hidden_states return_dictreturnc s"|dur |n|jj}|dur |n|jj}|dur4|n|jsB|jjnd}|durR|n|jj}|durltd|j|jkr| |durdurt dn|durdurt ddur| ||r|dur d|jj |jjffddtd D}|d d 8<|jr8|jr8|r8td d}} |rFd nd} |rTd nd} t|jD]t\} } | | |||d\} }}|jr|jjdkr| d|jjdkr| d} |r| | f} |rb| |f} qb|| } |r| | f} |stdd| || | fDSt| || | dS)a input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*): If passed along, the model uses the previous state in all the blocks (which will give the output for the `input_ids` provided as if the model add `state_input_ids + input_ids` as context). use_cache (`bool`, *optional*): If set to `True`, the last state is returned and can be used to quickly generate the next logits. NFz<`attention_mask` was passed, but it is unused in this model.zDYou cannot specify both input_ids and inputs_embeds at the same timez5You have to specify either input_ids or inputs_embedsrcs0g|](}tj|dkrjntjjdqS)rr)r>rEr0rFr1rrrrrrJsz%RwkvModel.forward..rgꌠ9Y>)FzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...r)rQrrrr2css|]}|dur|VqdSrr)rxrrrreurz$RwkvModel.forward..)rrQrr)rqrrtrainingruse_return_dictr%Z warning_oncer_rescale_layersr:rr9rSrrZr enumerater rescale_everyrrr)r~rrrrQrrrrrZall_self_attentionsZall_hidden_statesrblockrrrrrIsp          zRwkvModel.forwardcCs|j|j krdS|jjdkrtt8t|jD]\}}|jr|jj j dt ||jj|j jj dt ||jjq6t|jj j dr|jj j jdt ||jj|j jj jdt ||jjq6t|jj j dr||jj |||j j|q6|jj j dt ||jj|j jj dt ||jjq6Wdn1sj0Y|j |_dS)Nrr2SCB quant_state)rrrqrr>Zno_gradrrrrTrZmul_intrrPhasattrrdiv_ _bnb_4bit_dequantize_and_rescale)r~block_idrrrrr~s"  ""$ BzRwkvModel._rescale_layerscCsttstdddl}|j|jj|jj}|dt ||j j |j j |ddd|j}t|d|dS) z Perform the dequantization and rescaling of the weights of a given layer. After that operation the layer will be quantized again. z/Please install bitsandbytes to use this method.rNr2cpuF)Z requires_gradr)r ImportErrorZ bitsandbytesZ functionalZdequantize_4bitrrrrrrqrrZ Params4bitrKr1setattr)r~Z target_layerrZbnbZdequant_weightsZ quant_weightrrrrsz*RwkvModel._bnb_4bit_dequantize_and_rescale)NNNNNNNN)rVrWrXrprrr rr> LongTensorrrboolrrrrIrrrrrrrrs2   irz The RWKV Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). cseZdZdgZfddZddZddZdd d Zede e j e e j e e j e e e j e e j e ee ee ee eeeefd d d ZZS)RwkvForCausalLMz head.weightcs8t|t||_tj|j|jdd|_| dS)NFrm) rorprrrr|rSrheadrrrrrrps  zRwkvForCausalLM.__init__cCs|jSrrrrrrget_output_embeddingssz%RwkvForCausalLM.get_output_embeddingscCs ||_dSrrrrrrset_output_embeddingssz%RwkvForCausalLM.set_output_embeddingsNcKsT|dur|dddfd}|dur8|dur8d|i}nd|i}||d<||d<|S)NrlrrrQrr5)r~rrQrrkwargsZ model_inputsrrrprepare_inputs_for_generations z-RwkvForCausalLM.prepare_inputs_for_generation) rrrrQlabelsrrrrrc  Ks| dur | n|jj} |j||||||| d} | d} || } d}|durf|j| |fd|jji| }| s| f| dd}|dur|f|S|St|| | j| j| j dS)aJ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*): If passed along, the model uses the previous state in all the blocks (which will give the output for the `input_ids` provided as if the model add `state_input_ids + input_ids` as context). labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` use_cache (`bool`, *optional*): If set to `True`, the last state is returned and can be used to quickly generate the next logits. N)rrQrrrrrrr)rrrQrr) rqrrrZ loss_functionrrrQrr)r~rrrrQrrrrrrZ rwkv_outputsrrrrTrrrrIs@%  zRwkvForCausalLM.forward)NNN) NNNNNNNNN)rVrWrXZ_tied_weights_keysrprrrr rr>rrrrrrrrIrrrrrrs6    r)rrr)NF)NF)/rr dataclassesrpathlibrtypingrrr>Ztorch.utils.checkpointrZ generationrZmodeling_layersr Zmodeling_utilsr utilsr r r rrrZconfiguration_rwkvrZ get_loggerrVr%r#r*ZautogradFunctionr+rcrirrjrrrrrrr__all__rrrrsR          j , F&"T .l