g:PdZddlZddlmZmZmZmZmZmZddl Z ddl Z ddl m Z ddl m Z mZmZddlmZddlmZmZdd lmZdd lmZdd lmZdd lmZmZmZdd lm Z ddl!m"Z"ddl#m$Z$m%Z%m&Z&m'Z'm(Z(ddl)m*Z*m+Z+ddl,m-Z-e+rddl.m/Z/m0Z0ddl1m2Z2nd\Z2Z0Z/e*r ddl3m4Z4m5Z5nd\Z5Z4e6e2e0e4e5e/fZ7e'j8e9Z:dZ;Gdde j<Z=e"j>e=de j?de@de j?fdZAGdd eZBGd!d"e j<ZCGd#d$eCZDGd%d&eCZEeCeDeEd'ZFGd(d)e j<ZGGd*d+e j<ZHGd,d-e j<ZIGd.d/e j<ZJGd0d1e j<ZKd2ZLe$d3eLGd4d5e ZMd6ZNe$d3eLGd7d8eMZOGd9d:eMeZPe$d;eLGd<d=eMZQdS)>zPyTorch Zamba model.N)AnyDictListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache DynamicCache)GenerationMixin)AttentionMaskConverter)_flash_attention_forward)BaseModelOutputWithPastCausalLMOutputWithPast SequenceClassifierOutputWithPast)PreTrainedModel)ALL_LAYERNORM_LAYERS)add_start_docstrings%add_start_docstrings_to_model_forward#is_flash_attn_greater_or_equal_2_10loggingreplace_return_docstrings)is_causal_conv1d_availableis_mamba_ssm_available) ZambaConfig)mamba_inner_fnselective_scan_fn)selective_state_update)NNN)causal_conv1d_fncausal_conv1d_updateNNr!c,eZdZdfd ZdZdZxZS) ZambaRMSNormư>cttjt j||_||_dS)z; ZambaRMSNorm is equivalent to T5LayerNorm N)super__init__r Parametertorchonesweightvariance_epsilon)self hidden_sizeeps __class__s d/var/www/html/ai-engine/env/lib/python3.11/site-packages/transformers/models/zamba/modeling_zamba.pyr-zZambaRMSNorm.__init__QsD l5:k#:#:;; #c|j}|tj}|ddd}|tj||jzz}|j||zS)NT)keepdim) dtypetor/float32powmeanrsqrtr2r1)r3 hidden_states input_dtypevariances r7forwardzZambaRMSNorm.forwardYs|#) %((77  $$Q'',,R,>>% Ht?T4T(U(UU {]--k::::r8cHt|jjd|jS)Nz, eps=)tupler1shaper2r3s r7 extra_reprzZambaRMSNorm.extra_repr`s& )**II$2GIIIr8)r*)__name__ __module__ __qualname__r-rFrK __classcell__r6s@r7r)r)Psb$$$$$$;;;JJJJJJJr8r)rCn_repreturnc|j\}}}}|dkr|S|dddddddddf|||||}||||z||S)z This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) r N)rIexpandreshape)rCrQbatchnum_key_value_headsslenhead_dims r7 repeat_kvrZhs 2?1D.E h zz!!!!QQQaaa"23::5BUW\^bdlmmM  (;e(CT8 T TTr8c heZdZdZejdfdZ ddejdejdede e e e fde ejejff d Zd ejfd Zdde edefd Zde e eje ejffdZedde e e ejddfdZdS) HybridMambaAttentionDynamicCachea A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache (which has a constant shape regardless of seq_len). This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states` and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`, while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors). For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors), while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`, and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`. Nc ||_|j|_d|_|j|jz}|j}|j}|j|_g|_g|_ g|_ i|_ i|_ i|_ t|jD]}|xjt!j|||gz c_|j||jz|f} |xj t!j| |gz c_ |j|dkr|j |fdt|jD|_fdt|jD|_dS)NFdevicer=hybridcDg|]}tjggzSr_r/tensor.0_ batch_sizer_s r7 z=HybridMambaAttentionDynamicCache.__init__..s/rrrQ%,tj'8HHHrrrr8cDg|]}tjggzSrbrdrfs r7rjz=HybridMambaAttentionDynamicCache.__init__..s/tttqEL" ):6JJJtttr8)r=layers_block_typehas_previous_state mamba_expandr4 mamba_d_state mamba_d_conv n_mamba_heads conv_states ssm_statestransformer_layers_modules _parameters_buffersrangenum_hidden_layersr/zerosappend key_cache value_cache) r3configrir=r_intermediate_sizessm_state_sizeconv_kernel_sizei cache_shapes ` ` r7r-z)HybridMambaAttentionDynamicCache.__init__s !'!9"'"/&2DD-!.#1"$  v/00 2 2A    J(9;KTZbghhh!   &t'9;LPTPb;bdrsK OO Ke T T TU UOO%a(H44'..q111rrrrrRWX^XpRqRqrrrtttttTYZ`ZrTsTstttr8 key_states value_states layer_idx cache_kwargsrRcD|j|jddkr||j|<||j|<nVtj|j||gd|j|<tj|j||gd|j|<|j||j|fS)Nr;rr:dim)r|rIr}r/cat)r3rrrrs r7updatez'HybridMambaAttentionDynamicCache.updates >) $ *2 .! 3 3(2DN9 %*6D Y ' '(- 4>)3Lj2Y_`(a(a(aDN9 %*/)T5Ei5PR^4_ef*g*g*gD Y '~i($*:9*EEEr8beam_idxctt|jD];}|j|j}|j|d|||j|<|j|j}|j|d|||j|<|j|j}|j|d|||j|<|j|j}|j|d|||j|<=dS)zDReorders the cache for beam search, given the selected beam indices.rN) rxlenr|r_ index_selectr>r}rrrs)r3rrr_s r7 reorder_cachez.HybridMambaAttentionDynamicCache.reorder_cachesEs4>2233 i iI^I.5F(,y(A(N(NqRZR]R]^dReRe(f(fDN9 %%i07F*.*:9*E*R*RSTV^VaVabhViVi*j*jD Y '%i07F*.*:9*E*R*RSTV^VaVabhViVi*j*jD Y '_Y/6F)-)C)P)PQRT\T_T_`fTgTg)h)hDOI & & i ir8rc||jvr |jdn|}t|j|krdS|j|jdS)zYReturns the sequence length of the cached states. A layer index can be optionally passed.r)rtrr|rI)r3rs r7get_seq_lengthz/HybridMambaAttentionDynamicCache.get_seq_lengthsS3<4CZ2Z2ZD+A..`i t~  ) + +1~i(.r22r8c tdNzIHybridMambaAttentionDynamicCache does not have a legacy cache equivalent.NotImplementedErrorrJs r7to_legacy_cachez0HybridMambaAttentionDynamicCache.to_legacy_caches!"mnnnr8past_key_valuesrc tdrr)clsrs r7from_legacy_cachez2HybridMambaAttentionDynamicCache.from_legacy_caches""mnnnr8N)r)rLrMrN__doc__r/float16r-Tensorintrrstrrrr LongTensorrrr classmethod FloatTensorrr8r7r\r\ts  27tuuuu@26 FFLFlF F tCH~. F u|U\) * FFFF$ ie&6 i i i i33 3c3333ouU\':E%,@PW\]]] r8FrCattention_mask position_idspast_key_valueoutput_attentions use_cachecache_positionrRc  |\} } } ||} ||}||}| | | |j|jdd} || | |j|jdd}|| | |j|jdd}|| |||\}}t||j }t||j }tj | |ddtj|jdz z }|$|ddddddd|jdf}||z}t"j|dtj| j}t"j||j|j}tj ||}|| |j| |jfkr5t5d| |j| |jfd ||dd}|| | |j}|}||}|}|sd}|||fS) Nr r:r rr;)rr=)ptrainingz `attn_output` should be of size z , but is )sizerrrviewrrY transposerWrrZrr/matmulmathsqrtrIr functionalsoftmaxr?r>r=dropoutrrr contiguousrUrr)r3rCrrrrrrrkwargsbszq_lenrh query_statesrr attn_weights causal_mask attn_outputs r7rFzZambaAttention.forwards&**,, UA{{=11 [[// {{=11 #((eT^T]SS]]^_abcc __S%1I4=YYccdeghii #((eT5Mt}]]gghiklmm  %'5'<'>     #t~udm!L L L)CPTP]3^))$$&&))  "++Aq11<<>> !))#ud6PQQ ! kk+.. !   LL.88r8rNNNFFN)rLrMrNrr!rrr-r/rrr\boolrrFrOrPs@r7rrs,  ^^{^x}^^^^^^:2637EI"'598989|8989!. 89 u/0 89 !!AB 89 8989!!1289 u|Xel3XeEL>Q5RR S8989898989898989r8rceZdZdZfdZ ddejdedeejdeej d ee d e d e d eej fd Z xZ S)ZambaFlashAttention2aF Zamba flash attention module. This module inherits from `ZambaAttention` as the weights of the module stays untouched. The only required change would be on the forward pass where it needs to correctly call the public API of flash attention and deal with padding tokens in case the input contains any of them. cbtj|i|t |_dSr)r,r-r_flash_attn_uses_top_left_mask)r3argsrr6s r7r-zZambaFlashAttention2.__init__<s9$)&))) 3V2W2W.W+++r8NFrCrrrrrrrc |\} } } ||} ||}||}| | | |j|jdd} || | |j|jdd}|| | |j|jdd}|| |||\}}t||j }t||j }|j sdn|j }| j}|tjkrtjrtj}n3t'|jdr |jj}n|jjj}t.d|d| |} ||}||}| dd} |dd}|dd}dt5j|jdz z }t9| |||| ||}|| | |j}| |}|sd}|||fS)Nr r:_pre_quantization_dtypezThe input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in .)r softmax_scale)!rrrrrrrYrrWrrZrrrr=r/r?is_autocast_enabledget_autocast_gpu_dtypehasattrr~rr1logger warning_oncer>rrrrUrrr)r3rCrrrrrrrrrrrhrrr dropout_raterD target_dtyperrrs r7rFzZambaFlashAttention2.forwardDs&**,, UA{{=11 [[// {{=11 $((eT^T]SS]]^_abcc __S%1I4=YYccdeghii #((eT5Mt}]]gghiklmm  %'5'<'|<| } | } |}dt)j|j dz z }t,jj| | |||jr|jnd|jo |duo| dk|}| dd}|| | |j}||}|d|fS) NaZambaModel is using ZambaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.)rCrrrrrr r:rcudar) attn_mask dropout_prscale)rrr,rFrrrrrrrYrrWrrZrrIr_typerrrr/r rscaled_dot_product_attentionrrrrr)r3rCrrrrrrrrrrhrrrrrrr6s r7rFzZambaSdpaAttention.forwards     [   77??+-)-"3# # &**,, UA{{=11 [[// {{=11 #((eT^T]SS]]^_abcc __S%1I4=YYccdeghii #((eT5Mt}]]gghiklmm  %'5'<'> !&&sE43MNN kk+.. D.00r8r)rLrMrNrr/rrrrr\rrrFrOrPs@r7rrs 2637EI"'59G1G1|G1G1!. G1 u/0 G1 !!AB G1 G1G1!!12G1 u|Xel3XeEL>Q5RR SG1G1G1G1G1G1G1G1G1G1r8r)eagerflash_attention_2sdpacbeZdZdZdeffd Z d dejdefdZ d defdZ d defd Z xZ S) ZambaMambaMixeruE Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, and is why Mamba is called **selective** state spaces) This module differs from `transformers.models.mamba.modeling_mamba.MambaMixer` in two ways: - Added multi-head: the output of `self.in_proj` is split into `self.n_mamba_heads` heads, and each head undergoes an independent forward pass, identical to the original `MambaMixer`, up until the pre-activations of `self.out_proj`. The pre-activations, coming from different mamba heads, are then concatenated and fed into `self.out_proj`. r~c t||_||_|j|_|j|_|j|_|j |jz|_ |j |_ |j |_ |j |j z|_|j|_|j|_t'j|j |j |j|j|j |jdz |_|j|_t0|j|_|j|_t'j|j|j dz|j|_t'jt?j |j |j |jdzz|j|_!t'jt?j |j |j|j dz dz|j dzz |_"t'jt?j |j |j|_#t?j$d|jdzt>j%dddf}|&|j d'}t'jt?j(|)|j |jd|_*t'jt?j+|j |j|_,t'j|j |j|j|_-t\st^0ddSdS) Nr ) in_channels out_channelsr kernel_sizegroupspaddingr:rg?)r=r;apThe fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)` is None. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config)1r,r-r~rr4rorrprrnr mamba_dt_ranktime_step_rankrqmamba_head_dimmamba_conv_bias use_conv_biasmamba_proj_biasuse_biasr Conv1dconv1dhidden_mamba_act activationractuse_mamba_kernelsuse_fast_kernelsrin_projr.r/rz x_proj_weightdt_proj_weight dt_proj_biasaranger?rTrlogrUA_logr0Dout_projis_fast_path_availablerr)r3r~rAr6s r7r-zZambaMambaMixer.__init__s  "!-$2 & 3!'!4v7I!I$2#1"48JJ#3. i./#-))A-     !1&12 & 8y!143IA3MTXTabbb  \ &'$*=*AA'  !l [+T-@$BU V VY\ \ !3& '   LT5GI\)])]^^ LD/!35= I I I$PQPQPQ' R HHT+R 0 0 ; ; = =\%)A,,"6"6t7I4K^`b"c"cdd ej);T=PQQRR $"8$:JQUQ^___ %    ^       r8NrC cache_paramsc |j\}}}|duo |jo|dk}||dd}||dd|dd\}} |d}| d} | ||j d|dd} |j j |j j d|j j d} |r[t|d|j|j| |j j|j}|d}n|0t'j|dks||dz}|Zt*j||j|jdz df} |j|j| t5|| |j j|j}|0t'j|dks||dz}|d|j |j|dd}|jdddddddf|zdd} t'j| |j|j|jgd\} }}|j dddf| ddz}t'j!|j"# }|j$|j$#nd}t'j%|d|f|j&|j'}|rtQ|j D]}tS|j*|jdd|f||d df||d df||||dddf||dddf|j+|| |d df||d d}t'j,||fd}nOt'j%|d|j|jf|j&|j'}tQ|j D]}t[||||||||dd||dd|j+|#| |||d d \}}t'j,||fd}t'j,||dfd}|'|%|j*|j||.|dd}|S) Nr r:r;rr)rrr^.T) dt_softplus)delta_softplusreturn_last_state)/rIrmr rrchunksqueezerrUrqrr1rr&rrrrr unsqueezer/allr rpadrcopy_r%rr splitrrr exprfloatremptyr_r=rxr$rsrrr#r)r3rCrrriseq_lenrhuse_precomputed_statesprojected_statesgate conv_weightsrrssm_parameters time_stepBCdiscrete_time_steprtime_proj_bias scan_outputsn scan_outputs_ ssm_state ssm_state_contextualized_statess r7cuda_kernels_forwardz$ZambaMambaMixer.cuda_kernels_forward;s"/!4 GQ!-T!9!nl>]!nbimnbn << 66@@AFF.33JAwOOUUVW]^U__ t%--a00;;== ||A||J(:BHHRRSTVWXX{)..t{/A/F/Fq/I/I4;K]KbKbcdKeKeff ! L0%%b))(8   M*33B77MM)%)Na>{KKK,]L$+JZgkgvwwwM)%)Nar_rmrrr/rollsumrr1rrrrrr rrrrzrrr r!rr rsoftplusr"rr#rxrr{stackrr)r3 input_statesrrrir%rhr=r'rCr(rr3 conv_stater*r+r,r-r.r discrete_A discrete_BdeltaB_ur0r scan_outputr5s r7 slow_forwardzZambaMambaMixer.slow_forwards"!-!3 GQ"<< 55??1EE.33JAwOOUUVW]^U__ t%--a00;;== ||A||J(:BHHRRSTVWXX|-MNN ( L0@FqIZWW} D(3DNCIIKK (3DNC ! ]%9::I/ oqLL ,T^<B1ESS)5dnE "Z 2BGGG '4QQQ1W'= 111aaa8$;E (8 % *t{7I!!!QPQPQPQ'7R*RXZ [ [ [ %6!T[%55M $ 7 7 : :5 A A K KB O O !-eiRS@S6T6T-$1N111}GZ[]G^F^F`F`C`4a4k4klm4n4n$nM]..}t?TWdWjkmWn?npq>rss ;E (8 $])C)CC'M)R S S !-eiRS@S6T6T-$1N111}GZ[]G^F^F`F`C`4a4k4klm4n4n$nM T/1DdFYZ$+I )%)NaN ? E E G GG w 9 9A"111aaaAqqq=1;;AqAAIMPXYZYZYZ\]\]\]_`_`_`bcefefefYfPgPqPqrsuvPwPwwI,y':':1a'@'@'C'CE'J'JAaaaQRQRQRTUWXWXWXjMLcLcdfLgLghhK    AAAqqq!!!QJ 7 8 8 8 8k,B777 !]TVAAAtQQQ=1.2.0' and 'pip install mamba-ssm', or set use_mamba_kernels=False in the model's config.)r)r rr r_rrr6rF)r3rCrrs r7rFzZambaMambaMixer.forwardsz   i) V4;M;T;Y-Y-Y i ,,]LYg,hh h   ^ \\\r8r') rLrMrNrr!r-r/rr\r6rFrFrOrPs@r7rrs  ={======@rv_%_%"\_%9Y_%_%_%_%B[%[%7W[%[%[%[%z ] ]3S ] ] ] ] ] ] ] ]r8rc$eZdZfdZdZxZS)ZambaMLPct|j|_|j|_t j|j|jd|_t j|j|jd|_t j|j|jd|_t|j |_ dSNFr) r,r-r4rr r gate_projup_proj down_projr hidden_actact_fnr3r~r6s r7r-zZambaMLP.__init__s !-!'!94#3T5KRWXXXy!143IPUVVV 4#94;KRWXXXV./ r8c||||||zSr)rNrPrLrM)r3 hidden_states r7rFzZambaMLP.forwards>~~dkk$..*F*FGG$,,WcJdJddeeer8)rLrMrNr-rFrOrPs@r7rIrIsN00000fffffffr8rIc8eZdZddedeeffd Z ddejdejdedeejd eej d ee d ee d ee d eej de ej ee ej ej fffdZxZS)ZambaAttentionDecoderLayerNr~rc0tt|j|||_t ||_t|j|j |_ t|j |j |_ dS)Nr5) r,r-ZAMBA_ATTENTION_CLASSES_attn_implementation self_attnrI feed_forwardr)r rms_norm_epsinput_layernormr4pre_ff_layernormrs r7r-z#ZambaAttentionDecoderLayer.__init__s 01LMfV_``$V,,+F,HfNabbb ,V-?VEX Y Y Yr8FrCoriginal_hidden_statesrrrrrrrRc tj||gd}||}|jd|||||||| d| \}} } ||}||}|f} |r| | fz } |r| | fz } | S)a Args: hidden_states (`torch.FloatTensor`): output of previous Mamba layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output of shape `(batch, seq_len, embed_dim)`. This is concatenated with `hidden_states` (which is the output of the previous (mamba) layer). The concatenated tensor is then used as input of the pre-attention RMSNorm (see fig. 2 in https://arxiv.org/pdf/2405.16712). attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. r;r)rCrrrrrrrr)r/ concatenater]rZr^r[)r3rCr_rrrrrrrrself_attn_weightspresent_key_valueoutputss r7rFz"ZambaAttentionDecoderLayer.forwards@)=:P*QWYZZZ ,,];; >Ldn ? ')%)/) ? ?  ? ? ; (*;--m<< ))-88  "  , )+ +G  , )+ +Gr8rr)rLrMrNr!rrr-r/rrr\rrrrFrOrPs@r7rUrUs7ZZ{Zx}ZZZZZZ2637EI,1$)5999|9!& 9 9 !. 9 u/0 9!!AB9$D>9D>9!!129 u (51BEDU1U+V"WW X99999999r8rUcreZdZdedeffd Z ddejdeejdedeejd eejd eej d ee d ee d ee deej deejde ej ee ej ej fffdZxZS)ZambaMambaDecoderLayerr~rctt|||_t |j|j|_||_dS)N)r~rrW) r,r-rmambar)r4r\r]rrs r7r-zZambaMambaDecoderLayer.__init__ZsS $FiHHH +F,>FDWXXX"r8NFrCr_rrrrrrrtransformer_hidden_statesrRc |} | || zn|}||}||||}d} | |z}|f}|r|| fz }| r||fz }|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. N)rCrr)r]rh)r3rCr_rrrrrrrrriresidualrbrds r7rFzZambaMambaDecoderLayer.forward`s<! :S9^M5 5 5dq ,,];;  '')#  !!=0  "  , )+ +G  ) ( (Gr8) NNNNNNFFNN)rLrMrNr!rr-r/rrrr\rrrrFrOrPs@r7rfrfYsX#{#s######:>15.237EI,1$)59<@::|:!) 6: : !. : el+ :u/0:!!AB:$D>:D>:!!12:$,EL#9: u (51BEDU1U+V"WW X::::::::r8rfcdeZdZdedejdeffd Z ddej de ej d e d e ej d e ej d e ej d e e de ede ede ej deeje eejejfffdZxZS) HybridLayer shared_transflinearrhcrt||_||_||_dSr)r,r-rnro mamba_decoder)r3rnrorhr6s r7r-zHybridLayer.__init__s6 * "r8NFrCr_rrrrrrrrrRc ||||||||| |  } | d} |r| d} || } ||| ||||| | } |r| d| f| ddz} | S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with hidden activations to form the input of the shared transformer layer. layer_idx (`int`): layer number. attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. )r_rrrrrrrrr )rirrrrrrr:N)rnrorq)r3rCr_rrrrrrrr layer_outputsrirbs r7rFzHybridLayer.forwards@** #9&%)/)+   %2!$4!  1 -a 0 $(KK0I$J$J!** &?)%)/)+     V*1-/@AMRSRTRTDUUMr8) NNNNNNFFN)rLrMrNrUr rrfr-r/rrrrr\rrrrFrOrPs@r7rmrmsc#&@#")#\r######:>15.237EI,1$)59AA|A!) 6A A !. A el+ Au/0A!!ABA$D>AD>A!!12A u (51BEDU1U+V"WW XAAAAAAAAr8rmaJ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`ZambaConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. zSThe bare Zamba Model outputting raw hidden-states without any specific head on top.c eZdZeZdZdZddgZdZdZ dZ dZ dZ dZ ee dd eejd eeeeeeffd ed effd ZxZS)ZambaPreTrainedModelmodelTrUrfrFc|jj}t|tjtjfrJ|jjd||j |j j dSdSt|tj rU|jjd||j +|jj|j  dSdSt|trd|j_d|j_|jjd||jjdz}tj|j| ||jj|jjz|jjz}t3jt3j|jj|t9j|jjt9j|jjz zt9j|jjz |jj!}|t3jt3j"|  z}t3j#5|j$%|dddn #1swxYwYd|j$_&dSdS)Nr)rAstdTg)min)'r~initializer_ranger:r rrr1datanormal_rzero_ Embedding padding_idxrr_no_weight_decayrr rinituniform_r rnr4rqr/r"randrr time_step_max time_step_minclamptime_step_floorexpm1no_gradrr  _no_reinit)r3modulerx dt_init_stdrdtinv_dts r7 _init_weightsz"ZambaPreTrainedModel._init_weightssk+ fry")4 5 5 2 M  & &CS & 9 9 9{&  &&((((('&  - - 2 M  & &CS & 9 9 9!- "6#56<<>>>>>.-  0 0 2,0FL )(,FH %  % - -3C - @ @ @+3T9K G  V2[L+ N N N![5 8OOSWS^SllN 4;4nEE8DK566$+B[9\9\\^(4;4556e 3e44  %)U["%5%5$5666F 2 2#))&111 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2-1F  * * *' 2 2sJ((J,/J,N torch_dtype device_maphard_check_onlycheck_device_mapc|t|||||}|s|jdkrd|_|S)z Overloads `PreTrainedModel._check_and_enable_flash_attn_2` so as to DISABLE Flash Attention 2 by default on Zamba models. Flash attention 2 is currently not supported in the HuggingFace implementation of Zamba v1. )rrrr)r,_check_and_enable_flash_attn_2rY)rr~rrrrr6s r7rz3ZambaPreTrainedModel._check_and_enable_flash_attn_2'sT77 K__o8    26#>BU#U#U*1F ' r8)NNFF)rLrMrNr! config_classbase_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_2_supports_sdpa_supports_cache_class _is_statefulrrrr/r=rrrrrrrOrPs@r7rurus L&*#57OP"3"N L222>.2;? %!& ek*U3S#X#678     [[r8rua Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. If `past_key_values` is used, optionally only the last `input_ids` have to be input (see `past_key_values`). If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) past_key_values (`HybridMambaAttentionDynamicCache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): A HybridMambaAttentionDynamicCache object containing pre-computed hidden-states (keys and values in the self-attention blocks and convolution and ssm states in the mamba blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. Key and value cache tensors have shape `(batch_size, num_heads, seq_len, head_dim)`. Convolution and ssm states tensors have shape `(batch_size, d_inner, d_conv)` and `(batch_size, d_inner, d_state)` respectively. See the `HybridMambaAttentionDynamicCache` class for more details. If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, this tensor is not affected by padding. It is used to update the cache in the correct position and to infer the complete sequence length. cTeZdZdZdeffd ZdZdZee  dde j de e j d e e j d e ed e e jd e ed e ede ede ede e j deeeffdZdZxZS) ZambaModelz Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`ZambaDecoderLayer`] Args: config: ZambaConfig r~c  t||j|_|j|_t j|j|j|j|_t|}g}g}|j |_ t|j D]}|j |dkr%| t||8|j |dkrb| t j|jj|jjd| t||t#|}t#|}g}g|_t'|j D]\}}|dkred|d gd} g|j fd | D|_| t)|t+|t+|p| t+|t j||_|j|_t3|j|j |_d|_|dS) Nrh)rr`Frzlayers.r) z%shared_transf.self_attn.q_proj.weightz%shared_transf.self_attn.k_proj.weightz%shared_transf.self_attn.v_proj.weightz%shared_transf.self_attn.o_proj.weightz+shared_transf.feed_forward.gate_proj.weightz)shared_transf.feed_forward.up_proj.weightz+shared_transf.feed_forward.down_proj.weightz$shared_transf.input_layernorm.weightz%shared_transf.pre_ff_layernorm.weightcg|]}|zSrr)rgkey prefix_names r7rjz'ZambaModel.__init__..s FnFnFn]`{UXGXFnFnFnr8rW)r,r- pad_token_idr vocab_sizer r~r4 embed_tokensrUrlrxryr{rfrr~iter_tied_weights_keys enumeratermnext ModuleListlayersrYr)r\final_layernormgradient_checkpointing post_init) r3r~block mamba_layers linear_layersrrlayer_id layer_type tied_keysrr6s @r7r-zZambaModel.__init__sp    !. +L):FFDWXXX&+# r8c|jSrrrJs r7get_input_embeddingszZambaModel.get_input_embeddingss   r8c||_dSrrr3values r7set_input_embeddingszZambaModel.set_input_embeddingss!r8N input_idsrrr inputs_embedsrroutput_hidden_states return_dictrrRc >||n |jj}||n |jj}||n |jj}| | n |jj} |du|duz rt d|jr%|jr|rt dd}|| |}|} tj |} |r|t d| &tj | jd| j} || d}|||| } |rdnd}|rdnd}t%|jD]s\}}|r|| fz }|jr,|jr%||j| | ||| ||||| }n|| | ||| ||||| }|d} |r|d ||dfz }t|| } |r|| fz }|r|jsd |_|sdn|}| st1d | |||fDSt3| ||| S) NzaYou cannot specify both input_ids and inputs_embeds at the same time, and must specify either onezX`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.FzZamba requires an initialized `HybridMambaAttentionDynamicCache` to return a cache. None was provided, so no cache will be returned.r rcrr) r_rrrrrrrrTc3K|]}||V dSrr)rgvs r7 z%ZambaModel.forward..1s(ttqfgfsfsfsfsfsttr8)last_hidden_staterrC attentions)r~rrruse_return_dictrrrrrrr/r;rrIr_r_update_causal_maskrr_gradient_checkpointing_func__call__rrmrHr)r3rrrrrrrrrrrCr_rall_hidden_statesall_self_attnsrlayerrs next_caches r7rFzZambaModel.forwards'2C1N--TXT_Tq$8$D $+Jj "+!6IIDKq $D>q 'tnq d^q !!12q  u-- .q q q CBq h!!!!!!!r8rceZdZdeffd ZdZdZdZdZdZ dZ e e e ee  dd ejd eejdeejdeedeejdeejdeedeedeedeedeejdedeeeffdZ ddZxZS)ZambaForCausalLMr~c&t|t||_dg|jj|_|j|_t j|j|jd|_ | dS)Nzlm_head.weightFr) r,r-rrvrrr rr4lm_headrrQs r7r-zZambaForCausalLM.__init__as    '' #3"Tdj6S"T +y!3V5FUSSS  r8c|jjSrrvrrJs r7rz%ZambaForCausalLM.get_input_embeddingsk z&&r8c||j_dSrrrs r7rz%ZambaForCausalLM.set_input_embeddingsn"' r8c|jSrrrJs r7get_output_embeddingsz&ZambaForCausalLM.get_output_embeddingsqs |r8c||_dSrr)r3new_embeddingss r7set_output_embeddingsz&ZambaForCausalLM.set_output_embeddingsts % r8c||_dSrrv)r3decoders r7 set_decoderzZambaForCausalLM.set_decoderws  r8c|jSrrrJs r7 get_decoderzZambaForCausalLM.get_decoderzs zr8) output_typerNrrrrrrlabelsrrrrrnum_logits_to_keeprRc ||n |jj}| | n |jj} | | n |jj} ||||||||| | |  }|d}||dd| dddf}d}||j|||jfi| }| s|f|ddz}||f|zn|St|||j |j |j S)a Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. num_logits_to_keep (`int` or `None`, *optional*): Calculate logits for the last `num_logits_to_keep` tokens. If `None`, calculate logits for all `input_ids`. Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences. Returns: Example: ```python >>> from transformers import AutoTokenizer, ZambaForCausalLM >>> model = ZambaForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1") >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba-7B-v1") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```N) rrrrrrrrrrrr losslogitsrrCr) r~rrrrvr loss_functionrrrrCr)r3rrrrrrrrrrrr loss_kwargsrdrCrroutputs r7rFzZambaForCausalLM.forward}sMb2C1N--TXT_Tq%9$D $+Jj &1%q*A)A)C)C&CD #~';A'>>>%aaa&78 > Y_Q/tz$+O  %,*>)..0077;;a?L  % %n&91 = = =  F+AAA 0B/B/D/D,DE   $ $+];LL')=)=)?)?@L ,#2&"0&*k&D"0    r8) NNNNNNNNNNNr)NNNNNT)rLrMrNr!r-rrrrrrrrrr_CONFIG_FOR_DOCr/rrrr\rrrrrrFrrOrPs@r7rr`s{'''(((&&&+*+ABB+AP_```'+1537FJ59-1$(,0/3&*59"#V V #V !.V u/0 V ""BC V   12 V )*V D>V $D>V 'tnV d^V !!12V  V  u,, -V V V a`CBV v44444444r8ra The Zamba Model with a sequence classification head on top (linear layer). [`ZambaForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). cveZdZfdZdZdZee ddee j dee j dee j dee e ee jfd ee jd ee j d eed eed eedeede eeffdZxZS)ZambaForSequenceClassificationc t||j|_t||_|jj|_t j|j|jd|_ | dSrK) r,r- num_labelsrrvrr rr4scorerrQs r7r-z'ZambaForSequenceClassification.__init__sv     +'' "&*"?Yv14?OOO  r8c|jjSrrrJs r7rz3ZambaForSequenceClassification.get_input_embeddings(rr8c||j_dSrrrs r7rz3ZambaForSequenceClassification.set_input_embeddings+rr8NrrrrrrrrrrrRc | | n |jj} ||||||||| |  } | d} || } ||jd}n |jd}|jj|dkrt d|jjd}nv|rtj||jj ddz }||jdz}| | j }nd}| tj || j |f}d}|t| | j }|jjf|jdkr d|j_nN|jdkr7|jtjks|jtj kr d |j_n d |j_|jjdkrWt%}|jdkr1|||}n|||}n|jjd krGt)}||d|j|d}n*|jjd krt-}|||}| s|f| ddz}||f|zn|St/||| j| j| j S) a labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). N)rrrrrrrrrr z=Cannot handle batch sizes > 1 if no padding token is defined.r;rc regressionsingle_label_classificationmulti_label_classificationr)r~rrvr rIrrr/rrargmaxr>r_r problem_typer r=rr rr rr rrrCr)r3rrrrrrrrrrtransformer_outputsrCrrisequence_lengths pooled_logitsrloss_fctrs r7rFz&ZambaForSequenceClassification.forward.s(&1%X $D>X 'tnX d^X  u66 7X X X CBX X X X X r8r )Rrrtypingrrrrrrr/torch.utils.checkpointr torch.nnr r r activationsr cache_utilsrr generationrmodeling_attn_mask_utilsrmodeling_flash_attention_utilsrmodeling_outputsrrrmodeling_utilsr pytorch_utilsrutilsrrrrrutils.import_utilsrrconfiguration_zambar!&mamba_ssm.ops.selective_scan_interfacer"r#+mamba_ssm.ops.triton.selective_state_updater$ causal_conv1dr%r&rr get_loggerrLrrModuler)r{rrrZr\rrrrXrrIrUrfrmZAMBA_START_DOCSTRINGrurrrr rr8r7r-sl( :::::::::::::::: AAAAAAAAAA!!!!!!........))))))GFFFFF .-----111111-,,,,,QXXXXXXXXRRRRRRR@P=-~8DDDDDDDDD-7**.0@BVXfg  H % %JJJJJ29JJJ(L))) UU\ U# U%, U U U UXoXoXoXoXo|XoXoXov_9_9_9_9_9RY_9_9_9J^9^9^9^9^9>^9^9^9FN1N1N1N1N1N1N1N1d- S]S]S]S]S]biS]S]S]n f f f f fry f f fBBBBBBBBJAAAAARYAAAHHHHHH")HHHV"Y@@@@@?@@ @FBJYSSSSS%SS Snkkkkk+_kkk\ j j j j j %9j j j j j r8