g@dZddlZddlmZddlmZmZmZddlZddl Zddlm Z ddl m Z ddl mZdd lmZdd lmZdd lmZmZmZmZmZdd lmZmZd dlmZejeZ erddl!m"Z"ddl#m$Z$m%Z%ndZ"er ddl&m'Z'm(Z(nd\Z(Z'e)e"e'e(fZ*dZ+dZ,dej-de.fdZ/dZ0dZ1GddZ2Gddej j3Z4Gdde j3Z5Gd d!e j3Z6Gd"d#e j3Z7Gd$d%eZ8eGd&d'eZ9eGd(d)eZ:d*Z;d+Z<ed,e;Gd-d.e8Z=ed/e;Gd0d1e8eZ>dS)2zPyTorch MAMBA2 model.N) dataclass)OptionalTupleUnion)nn)CrossEntropyLoss)ACT2FN)GenerationMixin)PreTrainedModel) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardlogging)is_causal_conv1d_availableis_mamba_2_ssm_available) Mamba2Config)selective_state_update)mamba_chunk_scan_combined mamba_split_conv1d_scan_combined)causal_conv1d_fncausal_conv1d_update)NNz!mistralai/mamba-codestral-7B-v0.1r input_tensorpad_sizect|jdkr ddddd|ddfnddd|ddf}tjj||ddS)z Padding x tensor with `pad_size` on the seq_len dim (dim=1) Assumes that we only have tensors of either size 4 or 3 rconstant)modevalue)lenshapetorchr functionalpad)rr pad_shapes f/var/www/html/ai-engine/env/lib/python3.11/site-packages/transformers/models/mamba2/modeling_mamba2.pypad_tensor_by_sizer)?sj 47|7I3J3Ja3O3OAq!Q!Q//VWYZ\]_gijlmUnI 8  " "<ST " U UUc"t||}t|jdkr.||jdd||jdS||jdd||jd|jdS)z Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and simultaneously splitting it into chunk sequences. Assumes that we only have tensors of either size 4 or 3 r r)r)r"r#reshape)rr chunk_sizes r(reshape_into_chunksr0Js&lH==L < !####L$6q$92z [batch_size, intermediate_size, conv_kernel_size] ssm_states: Dict[int, torch.Tensor] # layer_idx -> [batch_size, intermediate_size, ssm_state_size] Nconfig batch_sizer5r4cnd__j_t jjz_fdtj D_ fdtj D_ j _ tj _dS)Nrc ~i|]9}|tjjdjzjzzj:S)r-r3)r$zerosintermediate_sizen_groups state_sizeconv_kernel_size).0irHrGr4r5selfs r( z(Mamba2Cache.__init__..sg     u{&V_).sU    u{F,fov?PY_gl   r*) seqlen_offsetr5 conv_kernelrOintr; hidden_sizerLrangenum_hidden_layers conv_states ssm_states hidden_act activationr act)rRrGrHr5r4s`````r(__init__zMamba2Cache.__init__s & 2!$V]V5G%G!H!H         6344           6344     !+&+,r* layer_idxnew_conv_statecache_positionreturncP|j|}|d|jdz }|dd}||j|dddd|f<|j||j|xx|z cc<|j|S)Nrrr,shiftsdims)r]clamprOrolltor4zero_)rRrcrdre conv_states r(update_conv_statezMamba2Cache.update_conv_states%i0 '--a1F1JKK__BR_88 +9+<+cttjt j||_||_dSrrsuperrbr Parameterr$r=weightvariance_epsilonrRrZeps __class__s r(rbzMambaRMSNormGated.__init__sB l5:k#:#:;; #r*Nc|j}|tj}|?|tj|tjz}|ddd}|tj ||j zz}|j ||zSNr-r,T)keepdim) r5rmr$float32rr%silupowmeanrsqrtrr)rR hidden_statesgate input_dtypevariances r(forwardzMambaRMSNormGated.forwards#) %((77  )BM,>,>twwu}?U?U,V,VVM $$Q'',,R,>>% Ht?T4T(U(UU {]--k::::r*rrrrurvrwrbr __classcell__rs@r(rrsQ$$$$$$ ; ; ; ; ; ; ; ;r*rc 0eZdZdZdedeffd Z d dejde e de ej d e ejfd Z d de e de ej d e ejfd Z d de e de ej d e ejfd ZxZS) Mamba2Mixeru 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) rGrcc t|j|_|j|_|j|_|j|_t|j |jz|_ t|j |_ ||_ |j |_ |j|_t |j|_|j|_|j|_|j|_|j|_|j|_|j|_|j|_|j|_|j d|jz|jzz|_t7j|j|j|j |j|j|jdz |_|j |jz|jz}t7j|j||j|_ t7j!tEj#|j|_$tEj%d|jdz}t7j!tEj&||_'d|j'_(tS|j |j|_*t7j!tEj#|j|_+d|j+_(t7j|j |j|j|_,|j|_tZst\/ddSdS)Nr-r) in_channels out_channelsbias kernel_sizegroupspaddingrTraThe fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)` is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d)0rrbrUrZrNssm_state_sizerXrOrYr;rLtime_step_rankrc use_conv_biasr_r`r ralayer_norm_epsilonrms_normrMrVr/time_step_limit time_step_min time_step_maxconv_dimrConv1dconv1dLinearuse_biasin_projrr$r=dt_biasarangelogA_log_no_weight_decayrnormDout_projis_fast_path_availablelogger warning_once)rRrGrcprojection_sizeArs r(rbzMamba2Mixer.__init__s )!-$/ & 2!$V]T5E%E!F!F!&"788"#1 +&+,"(";    +%5#1#1.T]1BTEX1XX i %*=&*     04=@4>Qy       |EJt~$>$>??  LDNQ. / /\%)A,,// &* #%d&<$BYZZZ ej8899"& $"8$:JQWQ`aaa  %    >       r*Nr cache_paramsreattention_maskc |j\}}}|j|jz}d|jzd|jz|jzz|jz} ||jdkr||d} | jd| z dz} | | |j|j|jg} tj | | d\}}} }}t||j |j |jjd|jj|j}tj ||j||gd\}}}tj|j }|ddddfdddddfd|j|jtj}|dddddfdd|j}|jddddfd|j}|jddddfd|j}|||j|jd|jz}|||j|jd|jz}|||j|j}t9|j|j ||||||d|d }|||j|jz}||| }||ddddf}n|N|jddkr=|jddkr,|j }||dddddfz|}||}tj|j }|j!d t)d fkrind |j!i}|j"r|tG||jjd|jj|j|f|j|j$d|j|jj|jj%|jj|jj|j|jd dd |\}}ngtj ||j|j|jgd\} }}tL |jdvr]|'||(dd(ddddd|f}nstM|(dd|jjd|jj|j(ddddd|f}tj ||j||gd\}}}|N|jddkr=|jddkr,|j }||dddddfz|}tS|||d|j||||||jd||||jdf|j$|jddd|jdd|\}}|'|%|j|j *||||d}||| }||}|S)Nr-rrr,r8.r5T)zr dt_softplusgrAdt_limitF) rr/seq_idxr`rmsnorm_weight rmsnorm_epsoutproj_weight outproj_biasheaddimngroupsnorm_before_gatereturn_final_states)rswish)xrrr`)r/rrrrrr)+r#rMrrLrUrWrsqueezerr$splitrr]rcrrrr`exprfloatr;rVrmrrrviewrr^rrr5rtrainingrr/rrra transposercopy_)rRrrrerrHseq_len_groups_time_state_size d_to_removein_projected_statesd_mlpsplit_projection_dimrhidden_states_B_CdtBCrrrhidden_states_reshapedoutr5projected_statesdt_limit_kwargs ssm_state time_step scan_outputs r(cuda_kernels_forwardz Mamba2Mixer.cuda_kernels_forward sJ"/!4 GQ!%1D!D$001t}3DtGZ3ZZ]a]kk   # (BQ(F(F"&,,}/D/DQ/G/G"H"H (.r2[@QFE$)5$2H$-Y]Yg#h 05 >Az4=!'!* 2MNNAz4=!'!* 2MNNA%2%7%7 DNTXTa%b%b "2'7&    M*..z4>DM;YZZM IImT::M-- ..qqq$|z-Mamba2Mixer.torch_forward.. s)%z%z%z\]&9!Xt&W&W%z%z%zr*r)rr):r#r5rrrLrMrrUrrr^rcclonermr4rWr]r$rlndimrsumrrrrrarrr%r&rOrKrVrrrr;rsoftplusrkrrr. contiguousrbmmrrepeatr/r)permuter@rD zeros_likecatrr)3rR input_statesrrerrHrrr5rrrrrrrorrrrdAdBdBxr^ssm_states_reshaped C_reshapedyr D_residualA_cumsumLG_intermediateGM_intermediateMY_diag decay_statesB_decay_contractionstatesprevious_states decay_chunkstates_permutedresult new_statesstate_decay_outC_times_statesstate_decay_out_permutedY_offrcontextualized_statesrs3` @r( torch_forwardzMamba2Mixer.torch_forwards;!-!3 GQ" LL)=)=a)@)@AA!'+a$2H.HHAPTP]L]`d`sLssuyvDDIJJ(8(>(>t5 t~V\^)?) ) %1dM2  #$/?EEGGI! ]%9::I)A--)5dnE "Z 2BGGG ANASWXAXAX}QQQ111W'='=^k 111aaa8$(8>>zJJJ % *--8H8O*P*PSWS^SefgfgfgijlmlmlmfmSn*ntv w w w %6!T[%55M $ 7 7 : :5 A A!!!T3, O - 7 7! < < ]..!*]-@-DDaH (8>>zJJJ $])C)C)M)MaPQ)R)R S STUTUTUW_X_W_abababTb c !-.2Fq2IA2M2MR`RfghRilmRmRm)/E%2^AAAqqq$J5O%O$S$STY$Z$ZM T^T]D4=AAAy!((a$-HHA]Q&&**1733A *b))!!!T3,7AA''T\(9::BR!344B)11*gr4=YY__aaM *gD4GHHNNPPA *gr43FGGMMOOAAt~>BBAAt~>BBA'DO*CCtVH *-? x-X-XXJ*ByM9M]())B.A&{%z%z%z%zboqrtuwxay%z%z%z "M1a !Q1%%A|A2...H +a..))Aqqq!!!QQQaaa23a111dAAAqqq!!!8K6LLN""r"**Ay\AIIaAq!,D,DY,OON""r"**A l]111aaa:%>>CCAFFF!9hqqq!!!QQQ|&'>q!Q'J'J $#''++.Fy.QQEA *b$.$-HHAJA!||aaa'111aaa'( *gr22A$)A'7==iHHHii4(( !% knnU.C.C D D$$r*c@tr0d|jjjjvr|||||S|j}|G|jddkr6|jddkr%||dddddfz|}| ||||S)Ncudarr) rrrr4typerr5r#rmr)rRrrrerr5s r(rzMamba2Mixer.forwardRs " jf 0C0J0O&O&O,,]L.Zhii i#  %.*>q*AA*E*E.J^_`JadeJeJe*^AAAqqq$J-GGKKERRM!!-~~^^^r*NNN)rurvrwrxrrYrbr$r{rrFr|rrrrrs@r(rrs@|@@@@@@@J/35915 KK|K{+K!!12 K !. KKKK\x%x%8Mx%cklql|c}x%U]^c^jUkx%x%x%x%|/35915 __{+_!!12 _ !. ________r*rc&eZdZdfd ZdZxZS) Mamba2RMSNormrcttjt j||_||_dS)zM Mamba2RMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm Nrrs r(rbzMamba2RMSNorm.__init__dsD l5:k#:#:;; #r*c|j}|tj}|ddd}|tj||jzz}|j||zSr) r5rmr$rrrrrr)rRrrrs r(rzMamba2RMSNorm.forwardls|#) %((77  $$Q'',,R,>>% Ht?T4T(U(UU {]--k::::r*rrrs@r(rrcsL$$$$$$;;;;;;;r*rcreZdZfdZ ddeedeejdeejfdZ xZ S) Mamba2Blockct||_||_|j|_t |j|j|_t|||_ dS)Nrrc) rrbrGrcresidual_in_fp32rrZrrrmixer)rRrGrcrs r(rbzMamba2Block.__init__usd  " & 7!&"4&:STTT  9=== r*Nrrerc|}|||jjj}|jr|t j}|||||}||z}|S)Nrrrer)rrmrr5r%r$rr&)rRrrrerresiduals r(rzMamba2Block.forward}s! -"2"29I9O"2"P"PQQ   2{{5=11H  ^dr#  !=0 r*r) rurvrwrbrrFr$r|r{rrrs@r(r"r"ts>>>>>/35915 {+!!12  !. r*r"c.eZdZdZeZdZdgZdZdZ dZ dS)Mamba2PreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. backboner"TcRt|tr?d|j_d|j_t jt j|jj tj |jj tj |jj z ztj |jj z|jj}|t j t j|  z}t j5|j|dddn #1swxYwYd|j_t|t*jrB|j:t1|jdds$t*j|jnJt|t*jr0t*j|j|jj|jjr| D]\}}|dvrt*j!|tj"d t j5|tj"|jj#z}dddn #1swxYwYdSdS) zInitialize the weights.T)minN _no_reinitF)std)zout_proj.weight)a)$ isinstancerrrrr$rrandrGrUmathrrrrktime_step_floorexpm1no_gradrrr/rrrgetattrinitzeros_ Embeddingnormal_rinitializer_rangerescale_prenorm_residualnamed_parameterskaiming_uniform_sqrtr\)rRmodulerinv_dtnameps r( _init_weightsz#Mamba2PreTrainedModel._init_weightss fk * * -,0FL )(,FH % 4;0118DK566$+B[9\9\\^(4;4556e 3e44 %)U["%5%5$5666F - -$$V,,, - - - - - - - - - - - - - - -(,FN % fbi ( ( N{&v{L%@@0GNN6;///  - - N GOOFMt{/LO M M M ; / F"2244 F Fa... G,,Q$)A,,,???FFTYt{'DEEEFFFFFFFFFFFFFFF F F F Fs$D??EE+"JJ J N) rurvrwrxr config_classbase_model_prefix_no_split_modulessupports_gradient_checkpointing _is_statefulrGr}r*r(r+r+sR  L"&&*#L(F(F(F(F(Fr*r+ceZdZUdZdZeejed<dZ ee ed<dZ ee ejed<dS) Mamba2Outputa% Class for the MAMBA2 model outputs. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. cache_params (`Mamba2Cache`): 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`. Includes both the State space model state matrices after the selective scan, and the Convolutional states hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. Nlast_hidden_staterr) rurvrwrxrOrr$ FloatTensor__annotations__rrFrrr}r*r(rNrNsi$6:x 12999*.L(;'...80`, 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) 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. cache_params (`Mamba2Cache`, *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 `cache_params` is returned and can be used to quickly generate the next logits. 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. z`The bare MAMBA2 Model transformer outputting raw hidden-states without any specific head on top.cPeZdZfdZdZdZdZeee e e e  dde ejde ejd e ed e ed e ed e ed e ejde ejdeee ffdZxZS) Mamba2Modelcttjjj|_tjfdtj D|_ d|_ tjj |_||j|dS)Nc2g|]}t|S)r$)r")rPidxrGs r(rz(Mamba2Model.__init__..4s&$s$s$sC[3%G%G%G$s$s$sr*Fr)rrbrr< vocab_sizerZ embeddings ModuleListr[r\layersgradient_checkpointingrrnorm_f"_register_load_state_dict_pre_hook load_hook post_initrRrGrs `r(rbzMamba2Model.__init__0s    ,v'8&:LMMm$s$s$s$sSXY_YqSrSr$s$s$stt &+##F$6Fz&Mamba2Model.forward..s(ffqXYXeXeXeXeXeffr*)rOrr)rGryrrxuse_return_dict ValueErrorr\r_rFr:r4r5r$rrXr^_gradient_checkpointing_func__call__rWr#r`tuplerN) rRrvrwrrxryrzrerkwargsrall_hidden_states mixer_blocks r(rzMamba2Model.forwardHs&%9$D $+Jj "+!6IIZ^Zg=rT[=R=Rmr %0%r{rrrrrs@r(rWrW+s~      )))+*+BCC& $1548.2$(/3&*5915N N E,-N   01N {+ N D> N 'tn N d^N !!12N !.N  ul" #N N N DC N N N N N r*rWz The MAMBA2 Model transformer with a language modeling head on top (linear layer with weights not tied to the input embeddings). ceZdZgZfdZdZdZdZdZ dde e de e j d e e j fd Zeeeeee  dd e e j d e e jde e de e j de ede ede ede e j d e e j deeeffdZxZS)Mamba2ForCausalLMct|t||_t j|j|jd|_| dS)NFr) rrbrWr,rrrZr[lm_headrcrds r(rbzMamba2ForCausalLM.__init__s^    #F++ y!3V5FUSSS  r*c|jSrrrrss r(get_output_embeddingsz'Mamba2ForCausalLM.get_output_embeddingss |r*c||_dSrrrrps r(set_output_embeddingsz'Mamba2ForCausalLM.set_output_embeddingss % r*c4|jSrr)r,rnrss r(rnz&Mamba2ForCausalLM.get_input_embeddingss}11333r*c6|j|Srr)r,rrrps r(rrz&Mamba2ForCausalLM.set_input_embeddingss}11.AAAr*Nrrerc ||jd|jdz}n |jd}|r|td|ddkr%|dddfd}|dddfd}nstjd||j}tj|d||jdz |j} tj|| gd}d}|jd|krUtj|d||jdz |j} tj|| gd}||d|i} nd |i} | ||||d | S) Nrz`cache_position` should not be None as it should have been initialized in `model.generate`, you are responsible for passing in a valid `cache_position` if you are calling `prepare_inputs_for_generation` directly with `use_cache=True`rr,r2r|r8rwrv)rrrxre) r#rr$rr4r=r:rupdate) rRrvrwrxrrerrpast_len extended_mask model_inputss r(prepare_inputs_for_generationz/Mamba2ForCausalLM.prepare_inputs_for_generations  $$*1- 0BBHH q)H  $% e a 1$$%aaae,Y7 !/2!6y!A "'a)BR!S!S!S!& "''**H~7KA7N,NWeWl!!! "'NM+JPQ!R!R!R#   "X - - "J##A&&>3G3J(JSaShM#Y 'FANNNN  $)=+];LL'3L"0 ,&"0      r*rsrvrwlabelsryrzrxrfc ||n |jj}||||||||| } | d} || |jjj} d}||| j}| dddddf }|dddf }t}|| d| d| d}|s| f| ddz}||f|zn|St|| | j| jS)a 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]` N)rrwryrzrxrerr.r,r)rTrUrr)rGrr,rrmrr5rr4rrrr:rSrr)rRrvrwrrryrzrxrerrmamba2_outputsrrUrT shift_logits shift_labelsloss_fctoutputs r(rzMamba2ForCausalLM.forwards2&1%rrrrrs@r(rrs&&&444BBB .25915@@ {+ @ !!12 @!.@@@@D+*+BCC&($1559.2-1/3&*$(15157 7 E,-7   127 {+ 7 )* 7 'tn 7 d^7 D>7 !.7 !.7  u** +7 7 7 DC 7 7 7 7 7 r*r)?rxr5 dataclassesrtypingrrrr$torch.utils.checkpointrtorch.nnr activationsr generationr modeling_utilsr utilsr rrrrutils.import_utilsrrconfiguration_mamba2r get_loggerrur+mamba_ssm.ops.triton.selective_state_updater!mamba_ssm.ops.triton.ssd_combinedrr causal_conv1drrallrrrr{rYr)r0rDrFModulerrrr"r+rNrSMAMBA2_START_DOCSTRINGrrWrr}r*r(rs !!!!!!)))))))))) %%%%%%!!!!!!))))))------WVVVVVVV......  H % %"RRRRRRmmmmmmmmm!8DDDDDDDDD-7**46FH\]^^9  VU\VSVVVV   ((8 8 8 8 8 8 8 8 v;;;;;;;;$a_a_a_a_a_")a_a_a_H ;;;;;BI;;;"")84F4F4F4F4FO4F4F4Fn =====;== =0 =====;== =6 :fm m m m m 'm m m ` U U U U U -U U U U U r*