o
    iL                    @   s  d dl Z d dlZd dlmZ d dlmZmZmZmZ d dl	Z	d dl	m
Z
 d dlmZmZmZ ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZ ddlmZmZmZ ddlmZmZ ddl m!Z!m"Z" ddl#m$Z$ ddl%m&Z&m'Z' ddl(m)Z) ddl*m+Z+m,Z, ddl-m.Z. e, rd dl/m0Z0 d dl1m2Z2m3Z3 nd\Z0Z2Z3e+ rd dl4m5Z5m6Z6 nd\Z6Z5e'7e8Z9G dd de	j
j:Z;G dd de
j:Z<G dd dZ=G d d! d!e
j:Z>d"e	j?d#e@d$e	j?fd%d&ZA	'dTd(e
j:d)e	j?d*e	j?d+e	j?d,ee	j? d-eBd.eBfd/d0ZCd1d2 ZDdUd3d4ZEG d5d6 d6e
j:ZFd7e	j?d8e@fd9d:ZGd;d< ZHd=d> ZIeJe0e5e6fZKG d?d@ d@e
j:ZLG dAdB dBe
j:ZMG dCdD dDe
j:ZNG dEdF dFe
j:ZOG dGdH dHe
j:ZPG dIdJ dJe"ZQe&G dKdL dLeQZRG dMdN dNeQeZSe&dOdPG dQdR dReQZTg dSZUdS )V    N)cycle)AnyCallableOptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss   )ACT2FN)Cache)GenerationMixin)AttentionMaskConverter)FlashAttentionKwargs)BaseModelOutputWithPastCausalLMOutputWithPast SequenceClassifierOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringlogging)deprecate_kwarg)is_causal_conv1d_availableis_mamba_ssm_available   )Zamba2Config)selective_state_update)mamba_chunk_scan_combined mamba_split_conv1d_scan_combinedNNN)causal_conv1d_fncausal_conv1d_updateNNc                       s(   e Zd Zd fdd	ZdddZ  ZS )	Zamba2RMSNormGatedư>c                    s,   t    tt|| _|| _|| _d S N)	super__init__r   	Parametertorchonesweightvariance_epsilon
group_size)selfhidden_sizer1   eps	__class__ g/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/transformers/models/zamba2/modeling_zamba2.pyr+   =   s   

zZamba2RMSNormGated.__init__Nc           	      C   s   |j }|tj}|d ur|tj|tj }|j^ }}|| j }|j	g ||| jR  }|
djddd}|t|| j  }|j	g ||| j R  }| j|| S N   T)keepdim)dtypetor-   float32r   
functionalsilushaper1   viewpowmeanrsqrtr0   r/   )	r2   hidden_statesgateinput_dtypeprefix_dimslast_dimgroup_counthidden_states_groupvariancer7   r7   r8   forwardC   s   
zZamba2RMSNormGated.forwardr(   r)   )__name__
__module____qualname__r+   rO   __classcell__r7   r7   r5   r8   r'   <   s    r'   c                       s.   e Zd Zd fdd	Zdd Zdd Z  ZS )	Zamba2RMSNormr(   c                    s&   t    tt|| _|| _dS )z<
        Zamba2RMSNorm is equivalent to T5LayerNorm
        N)r*   r+   r   r,   r-   r.   r/   r0   )r2   r3   r4   r5   r7   r8   r+   R   s   

zZamba2RMSNorm.__init__c                 C   sJ   |j }|tj}|djddd}|t|| j  }| j|| S r9   )	r=   r>   r-   r?   rD   rE   rF   r0   r/   )r2   rG   rI   rN   r7   r7   r8   rO   Z   s
   zZamba2RMSNorm.forwardc                 C   s   t | jj d| j S )Nz, eps=)tupler/   rB   r0   r2   r7   r7   r8   
extra_repra   s   zZamba2RMSNorm.extra_reprrP   )rQ   rR   rS   r+   rO   rX   rT   r7   r7   r5   r8   rU   Q   s    rU   c                   @   s   e Zd ZdZdZejdfdededej	de
e fdd	Zd
d Zdedeejejf fddZ	d!dejdejdede
eeef  deejejf f
ddZdejfddZd"de
e defddZdedejdejdejfddZdd  ZdS )#Zamba2HybridDynamicCachea  
    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)`.
    FNconfig
batch_sizer=   devicec              	      s  || _ |j| _d| _t|j|j | _|j| _|j	| _
|j| _g | _i | _i | _i | _i | _i | _t|jD ]7}tj | jd|j |j  | j
|d| j|< tj | j|j| j|d| j|< | j| dkrm| j| q6 fddt|jD | _ fddt|jD | _d S )NFr:   r\   r=   hybridc                        g | ]}t jg g  d qS r\   r-   tensor.0_r[   r\   r7   r8   
<listcomp>        z5Zamba2HybridDynamicCache.__init__.<locals>.<listcomp>c                    r_   r`   rb   rd   rg   r7   r8   rh      ri   )r=   layers_block_typehas_previous_stateintmamba_expandr3   intermediate_sizemamba_d_statessm_state_sizemamba_d_convconv_kernel_sizen_mamba_headstransformer_layers_modules_parameters_buffersconv_states
ssm_statesrangenum_hidden_layersr-   zerosmamba_ngroupsmamba_headdimappend	key_cachevalue_cache)r2   rZ   r[   r=   r\   ir7   rg   r8   r+   u   s:    z!Zamba2HybridDynamicCache.__init__c                 C   s
   t | jS r)   )lenr   rW   r7   r7   r8   __len__   s   
z Zamba2HybridDynamicCache.__len__	layer_idxreturnc                 C   s   | j | | j| fS r)   )r   r   r2   r   r7   r7   r8   __getitem__   s   z$Zamba2HybridDynamicCache.__getitem__
key_statesvalue_statescache_kwargsc                 C   sz   | j | jd dkr|| j |< || j|< ntj| j | |gdd| j |< tj| j| |gdd| j|< | j | | j| fS )Nr;   r   r:   dim)r   rB   r   r-   cat)r2   r   r   r   r   r7   r7   r8   update   s   
zZamba2HybridDynamicCache.updatebeam_idxc                 C   s   t t| jD ]V}| j| j}| j| d||| j|< | j| j}| j| d||| j|< | j| j}| j| d||| j|< | j| j}| j| d||| j|< qdS )zDReorders the cache for beam search, given the selected beam indices.r   N)	rz   r   r   r\   index_selectr>   r   rx   ry   )r2   r   r   r\   r7   r7   r8   reorder_cache   s    z&Zamba2HybridDynamicCache.reorder_cacher   c                 C   sL   || j vr
| j d n|}t| j|ks| j|  dkrdS | j| jd S )zYReturns the sequence length of the cached states. A layer index can be optionally passed.r   )rt   r   r   numelrB   r   r7   r7   r8   get_seq_length   s    z'Zamba2HybridDynamicCache.get_seq_lengthnew_conv_statecache_positionc                 C   sr   | j | }|d| jd }|jddd}||j|d d d d |f< | j |   | j |  |7  < | j | S )Nr   r   r;   shiftsdims)rx   clamprr   rollr>   r\   zero_)r2   r   r   r   
conv_stater7   r7   r8   update_conv_state   s   

z*Zamba2HybridDynamicCache.update_conv_statec                 C   s   | j   | j  d S r)   )rx   r   ry   rW   r7   r7   r8   reset   s   
zZamba2HybridDynamicCache.resetr)   )r   )rQ   rR   rS   __doc__is_compileabler-   float16r   rl   r=   r   strr+   r   rV   Tensorr   dictr   r   
LongTensorr   r   r   r   r7   r7   r7   r8   rY   e   sN    
 

rY   c                       sD   e Zd ZU ejed< ddef fddZe e	dd Z
  ZS )	Zamba2RotaryEmbeddinginv_freqNrZ   c                    s   t    t|drt|jtr|jd|jd| _nd| _|j| _	|j| _
|| _t| j | _| | j|\}| _| jd|dd | j| _d S )Nrope_scaling	rope_typetypedefaultr   F)
persistent)r*   r+   hasattr
isinstancer   r   getr   max_position_embeddingsmax_seq_len_cachedoriginal_max_seq_lenrZ   r   rope_init_fnattention_scalingregister_bufferr   original_inv_freq)r2   rZ   r\   r   r5   r7   r8   r+      s   
zZamba2RotaryEmbedding.__init__c           
      C   s   | j d d d d f  |jd dd|j}|d d d d d f  }t|jjtr6|jjdkr6|jjnd}t	j
|dd+ | |  dd}t	j||fdd	}| | j }| | j }	W d    n1 smw   Y  |j|jd
|	j|jd
fS )Nr   r;   r   mpscpuF)device_typeenabledr:   r   r=   )r   floatexpandrB   r>   r\   r   r   r   r-   autocast	transposer   cosr   sinr=   )
r2   xposition_idsinv_freq_expandedposition_ids_expandedr   freqsembr   r   r7   r7   r8   rO      s   0&zZamba2RotaryEmbedding.forwardr)   )rQ   rR   rS   r-   r   __annotations__r   r+   no_gradr   rO   rT   r7   r7   r5   r8   r      s   
 
r   rG   n_repr   c                 C   s^   | j \}}}}|dkr| S | dddddddddf |||||} | ||| ||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)rB   r   reshape)rG   r   batchnum_key_value_headsslenhead_dimr7   r7   r8   	repeat_kv   s
   0r           modulequerykeyvalueattention_maskscalingdropoutc                 K   s   t || j}t || j}	t||dd| }
|d ur3|d d d d d d d |jd f }|
| }
tjj|
dtj	d
|j}
tjj|
|| jd}
t|
|	}|dd }||
fS )Nr:   r   r   r;   )r   r=   )ptrainingr   )r   num_key_value_groupsr-   matmulr   rB   r   r@   softmaxr?   r>   r=   r   r   
contiguous)r   r   r   r   r   r   r   kwargsr   r   attn_weightscausal_maskattn_outputr7   r7   r8   eager_attention_forward  s   
&r   c                 C   sH   | dd| j d d f }| d| j d d df }tj| |fddS )z*Rotates half the hidden dims of the input..Nr;   r:   r   )rB   r-   r   )r   x1x2r7   r7   r8   rotate_half  s   r   c                 C   sD   | |}| |}| | t| |  }|| t||  }||fS )a  Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    )	unsqueezer   )qkr   r   r   unsqueeze_dimq_embedk_embedr7   r7   r8   apply_rotary_pos_emb#  s
   

r   c                       s   e Zd ZdZ			ddedee dee dee f fddZed	d
dd			dde	j
dedee	j
 d
ee deee	j
e	j
f  dee dee	j
ee	j
 eee	j
  f fddZ  ZS )Zamba2AttentionaZ  
    Multi-headed attention from 'Attention Is All You Need' paper.

    Adapted from transformers.models.mistral.modeling_mistral.MistralAttention:
    The input dimension here is attention_hidden_size = 2 * hidden_size, and head_dim = attention_hidden_size // num_heads.
    The extra factor of 2 comes from the input being the concatenation of original_hidden_states with the output of the previous (mamba) layer
    (see fig. 2 in https://huggingface.co/papers/2405.16712).
    Additionally, replaced
    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) with
    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim/2)
    Finally, this attention layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this
    layer is tied, un-tied adapters (formally the same as LoRA but used in the base model) modules are added to the q, k, v projectors to increase
    expressivity with a small memory overhead (see Fig. 2 of https://huggingface.co/papers/2411.15242).
    NrZ   r   num_fwd_mem_blocksblock_idc           	   	      s  t    || _|| _|j| _|j| _|j|j | _	|j
| _
| jd d | _d| _|j| _tj|j|j| j dd| _tj|j|j| j dd| _tj|j|j| j dd| _tj|j| j |jdd| _|| _|j| _|| _|jrtg | _tg | _tg | _t| jD ]p}||j |krt tj| j| jj!ddtj| jj!| jdd}t tj| j| jj!ddtj| jj!| jdd}t tj| j| jj!ddtj| jj!| jdd}nt" }t" }t" }| j#| | j#| | j#| qdd t$| jD | _%d S )Nr:   g      TFbiasc                 S      i | ]\}}||qS r7   r7   re   indexr   r7   r7   r8   
<dictcomp>      z,Zamba2Attention.__init__.<locals>.<dictcomp>)&r*   r+   rZ   r   attention_hidden_sizeattention_head_dimr   num_attention_headsr   r   r   r   	is_causalattention_dropoutr   Linearq_projk_projv_projr3   o_projr   hybrid_layer_idslayer_block_mapr   use_shared_attention_adapter
ModuleListlinear_q_adapter_listlinear_k_adapter_listlinear_v_adapter_listrz   num_mem_blocks
Sequentialadapter_rankIdentityr   	enumerate	layer_dic)	r2   rZ   r   r   r   r   linear_q_adapterlinear_k_adapterlinear_v_adapterr5   r7   r8   r+   N  sT   
zZamba2Attention.__init__past_key_valuepast_key_values4.58new_nameversionrG   r   position_embeddingsr   r   c                 K   sp  |j d d }g |d| jR }| |}	| |}
| |}| jjrD| j| }|	| j| | }	|
| j	| | }
|| j
| | }|	|dd}	|
|dd}
||dd}| jjrp|\}}t|	|
||\}	}
|d ur}||
||\}
}t}| jjdkrt| jj }|| |	|
||f| jsdn| j| jd|\}}|jg |dR   }| |}||fS )Nr;   r   r:   eagerr   )r   r   )rB   r   r   r  r  rZ   r  r  r  r	  r
  rC   r   use_mem_roper   r   r   _attn_implementationr   r   r   r   r   r   r  )r2   rG   r   r   r  r  r   input_shapehidden_shapequery_statesr   r   adapter_layer_idxr   r   attention_interfacer   r   r7   r7   r8   rO     sH   






zZamba2Attention.forwardr#   )rQ   rR   rS   r   r   r   rl   r+   r   r-   r   rY   rV   r   r   rO   rT   r7   r7   r5   r8   r   >  sB    8r   input_tensorpad_sizec                 C   sH   t | jdkrddddd|ddfnddd|ddf}tjjj| |dddS )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
       r   constant)moder   )r   rB   r-   r   r@   pad)r#  r$  	pad_shaper7   r7   r8   pad_tensor_by_size  s   2r*  c                 C   sX   t | |} t| jdkr| | jd d|| jd S | | jd d|| jd | jd S )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*  r   rB   r   )r#  r$  
chunk_sizer7   r7   r8   reshape_into_chunks  s   
r,  c                 C   s   |  d}| d jg |   |R  } tjtj||| jtjddd}| | d} tj| dd}tjtj||| jtjddd}|| tj	 }|S )zo
    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
    r;   .Nr]   diagonalr   r   r   )
sizer   r-   trilr.   r\   boolmasked_fillcumsuminf)r#  r+  masktensor_segsumr7   r7   r8   segment_sum  s   
  r8  c                       s   e Zd ZdZddedee f fddZ		ddej	dee
 d	eej	 fd
dZddee
 d	eej	 fddZ		ddee
 d	eej	 fddZ  ZS )Zamba2MambaMixeru  
    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)
    NrZ   r   c                    s  t    || _|j| _|j| _|j| _t|j	| j | _
|| _|j| _d| _t | _|j| _|j| _|j| _| jj| _|j| _|j| _|j| _|j| _| j
d| j | j  | _tj| j| jd|j| j|jd d| _| j
| j | j }tj| j||j d| _!t"t#$| j| _%t#&d| jd }t"t#'|| _(t)| j
| j
| j dd| _*t"t#$| j| _+tj| j
| j|j d| _,t-st./d	 d S d S )
NrA   r:   Tr   )in_channelsout_channelsr   kernel_sizegroupspaddingr   gh㈵>)r1   r4   a  The fast path is not available because one 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)0r*   r+   rZ   r3   ro   rp   rq   rr   rl   rm   rn   r   use_conv_bias
activationr   SiLUactuse_mem_eff_pathr}   n_groupsr~   r   rs   	num_headsr+  time_step_limittime_step_mintime_step_maxconv_dimConv1dconv1dr   add_bias_linearin_projr,   r-   r.   dt_biasarangelogA_logr'   normDout_projis_fast_path_availableloggerwarning_once)r2   rZ   r   projection_sizeAr5   r7   r8   r+     s\   



	zZamba2MambaMixer.__init__rG   cache_paramsr   c                 C   sF  |j \}}}| j| j }d| j d| j | j  | j }|d ur|jr| |d}	|	j d | d }
|
|
| j| j| jg}t	j
|	|dd\}}}}}t||j| j | jjd| jj| j}t	j
|| j||gdd\}}}t	| j  }|d d d df d d d d d f d| j| jjt	jd}|d d d d d f dd| j}| jd d d df d| j}| jd d d df d| j}||| j|j d | j }||| j|j d | j }||| j| j}t|j| j ||||||d |dd
}||| j| j }| ||}| |d d d df }|S |d ur;t	 |dks;|j!}||d d d d d f  |}| |}t	| j  }| j"d u rQi nd	| j"i}|d urct	 |dk}nd}| j#r| j$r|d u r|rt%|| jjd| jj| j|f| j| j&d | j| jj| jj'| jj| jj| j| jd
dd|\}}|S t	j
|| j| j| jgdd\}}}|d ur|(dd}t)j*+|| j,|j d  df}|j| j -| t.d u s| jdvr| /| |(dd(ddd d d |f }n t.|(dd| jjd| jj| jd(ddd d d |f }t	j
|| j||gdd\}}}|d urNt	 |dksN|j!}||d d d d d f  |}t0|||d| j|||||| jd|||| jdf| j&| jd d d| jdd|\}}|d ur|d ur|j| j -| |||d}| ||}| |}|S )Nr:   r   r;   r   .r   T)zrN  dt_softplusdt_limitF)rS  r+  seq_idxr@  rmsnorm_weightrmsnorm_epsoutproj_weightoutproj_biasheaddimngroupsnorm_before_gatereturn_final_statesr   )rA   swish)r   r/   r   r@  )r+  rS  r[  r^  rf  rN  r\  )1rB   rD  rp   rn   rE  rk   rM  squeezerI  r-   splitr%   rx   r   rK  r/   r   r@  exprQ  r   r   r   r>   r?   rN  rS  rC   r    ry   rR  rT  allr=   rF  rC  r   r"   r+  r0   r   r   r@   r(  rr   copy_r$   rB  r!   )r2   rG   rZ  r   r[   seq_lenrf   groups_time_state_sized_to_removein_projected_statesd_mlpsplit_projection_dimrH   hidden_states_B_CdtBCrY  rN  rS  hidden_states_reshapedoutr=   projected_statesdt_limit_kwargsinput_not_masked	ssm_state	time_stephidden_states_B_C_tr   scan_outputr7   r7   r8   cuda_kernels_forward;  s   

<"
] 

 
L
(

 

z%Zamba2MambaMixer.cuda_kernels_forwardc           1   
      s	  |j \}}}|j}|d ur|jr|d}n |d ur4t|dks4||d d d d d f  |}|}|j d dj  dj	 j
  j d }	|j|	|	jjjgdd\}}}
}}|d ur8|jj  }||j}|jr|
d}
|jj }tj|ddd}|jdkr|d d dd d f n||d d d d df< |jj | tj||jjjd d dd d f  dd}jr|jj7 }||d d d df }n||dd}tj |j!|j d  df}|jj | |ddd d d |d d f }|d ur7t|dks7|j}||d d d d d f  |}n&tj"|jj#j
f|j|d	}|dddd |f dd}tj|jj	j
 j	j
 gdd\}}}t$j%&  }|d ur|jr|jdkr|d d d df n|d d dd d f d d d df }|dd'||j d j#}j(d
 'j(j d j#}tjj)|||j }t*|j+}|d 'jj#j
jtj,d}t$|d
 | }|-|j	ddd d d f }|'|j	jj	 |j d . }|-|d|j d }|d
 |dd d d f  }|-|dj#}||d
  }|jj |jj | |  |-|j	ddd d d f }|'|j	jj	 |j d . }|-|d|j d }|jj |j}|/|j j#j
}|/|j j
d}t0||}|/|jj#}j1d
 'j1j d j#}|||  |j}|-|dd d d df }ntj)|j( }t*|j+}|-||dj#& }|-||dj
& }|-||dj
& }|j2jj	 djd}|j2jj	 djd}j3|j3  j3  j1d
 t4|  }||d
  }||j| } fdd||||fD \}}}}|5dddd}tj6|dd}t$t7|}|d d d d d d d d d d d f |d d d d d d d d d d d f  }|jdd}|d
 |5dddddd
  } | jdd}!|!d
 |d d d d d f  d}"t$|d d d d d d dd f | }#||#5ddddd
  }$|$5dddddd
 |5ddddddd d d f  jdd5ddddd}%|d ur|jr|jj d d d df }&nt8|%d d d df }&tj9|&|%gdd}%t$t7tj |d d d d d d df d}'|%5ddddd}(|'d |(d d d d d df  jdd})|)5ddddd}*|*d d d df |*d d df }%}t$|}+|dd d d f |%d d d d d df  },|+5dddd}-|,d|-d
  }.|"|. }|-|djj#}|| } dkr|d d d |d d d d f }|-||d}|d ur|d ur|jj | :||
}/;|/|}0|0S )Nr   r;   r:   r   r   r   r   .r]   r-  ).NNr   )r   output_sizec                    s   g | ]	}t | jqS r7   )r,  r+  )re   tr$  r2   r7   r8   rh   M  s    z2Zamba2MambaMixer.torch_forward.<locals>.<listcomp>r%  )r   r   )<rB   r=   rk   rM  rh  r-   rk  r>   rn   rD  rp   rE  ri  rI  ry   r   cloner\   r   rx   r   ndimrl  sumrK  r/   r?  r   rB  r   r   r@   r(  rr   r|   r   rj  rQ  r   r   rN  softplusr   rG  r?   r   r   rC   bmmrS  repeat_interleaver+  r*  permuter4  r8  
zeros_liker   rR  rT  )1r2   input_statesrZ  r   r[   rm  rf   r=   ry  rq  rH   rG   rt  r|  r   ru  rv  rY  rN  dAdBdBxry   ssm_states_reshaped
C_reshapedyrS  
D_residualA_cumsumLG_intermediateGM_intermediateMY_diagdecay_statesB_decay_contractionstatesprevious_statesdecay_chunkstates_permutedresult
new_statesstate_decay_outC_times_statesstate_decay_out_permutedY_offr  contextualized_statesr7   r  r8   torch_forward  s    
.

60 . ,.B"$$$P$*L0(&
*
 zZamba2MambaMixer.torch_forwardc                 C   s0   t rd| jjjjv r| |||S | |||S )Ncuda)rU  rM  r/   r\   r   r  r  )r2   rG   rZ  r   r7   r7   r8   rO     s   zZamba2MambaMixer.forwardr)   r&   )rQ   rR   rS   r   r   r   rl   r+   r-   r   rY   r  r  rO   rT   r7   r7   r5   r8   r9    s,    B
  Fr9  c                       s6   e Zd Zddedee f fddZd	ddZ  ZS )
	Zamba2MLPNrZ   r   c              	      s   t    || _|j| _|j| _|| _|| _tj| jd| j |j	d| _
tj| j| j|j	d| _t|j | _tg | _t| jD ]/}||j |krfttj| jj| jjddtj| jjd| j dd}nt }| j| qA|j}dd t|D | _dS )aQ  
        This MLP layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this layer
        is tied, un-tied adapter modules (formally same as LoRA, but used in the base model) are added to the up and gate projectors to increase expressivity with a small memory overhead.
        r:   r   Fc                 S   r   r7   r7   r   r7   r7   r8   r     r   z&Zamba2MLP.__init__.<locals>.<dictcomp>N)r*   r+   rZ   r3   rn   r   r   r   r   rL  gate_up_proj	down_projr   
hidden_actact_fnr  gate_up_proj_adapter_listrz   r  r  r  r  r   r  r  r  )r2   rZ   r   r   r   gate_up_proj_adapterr  r5   r7   r8   r+     s(   
zZamba2MLP.__init__c                 C   sZ   |  |}| j| }|| j| | }tj|ddd}| |d |d  }| |}|S )Nr:   r;   r   r   r   )r  r  r  r-   chunkr  r  )r2   hidden_stater   gate_up_stateoutputr7   r7   r8   rO     s   


zZamba2MLP.forwardr&   r)   )	rQ   rR   rS   r   r   rl   r+   rO   rT   r7   r7   r5   r8   r    s    r  c                       s   e Zd Zddedee dee f fddZeddd	d
				ddej	dej	dedeej	 dee
 dee deej dee deejeeejejf  f fddZ  ZS )Zamba2AttentionDecoderLayerNrZ   r   r   c                    sd   t    || _t|j}t|d||d| _t|||d| _t	|j
|jd| _t	|j|jd| _d S )Nr;   )r   r   r   )r   r   r4   )r*   r+   r   r   r  r   	self_attnr  feed_forwardrU   r   rms_norm_epsinput_layernormr3   pre_ff_layernorm)r2   rZ   r   r   num_gsr5   r7   r8   r+     s   

z$Zamba2AttentionDecoderLayer.__init__r  r  r  r  FrG   original_hidden_statesr   output_attentionsr  r   r   c              	   K   sl   t j||gdd}| |}| jd||||||d|\}}	| |}| ||}|f}
|r4|
|	f7 }
|
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://huggingface.co/papers/2405.16712).
            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
                `(batch, sequence_length)` where padding elements are indicated by 0.
            past_key_values (`Zamba2HybridDynamicCache`, *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`).
            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
                with `head_dim` being the embedding dimension of each attention head.
        r;   r   )rG   r   r   r  r  r  Nr7   )r-   concatenater  r  r  r  )r2   rG   r  r   r   r  r  r  r   self_attn_weightsoutputsr7   r7   r8   rO     s$    




z#Zamba2AttentionDecoderLayer.forwardr&   )NNFN)rQ   rR   rS   r   r   rl   r+   r   r-   r   rY   r2  r   r   r   rV   FloatTensorrO   rT   r7   r7   r5   r8   r    s4    $		
r  c                       s   e Zd Zdedef fddZedddd											
	
				ddejde	ej de	e de	ej de	ej de	e
 de	e de	e de	ej de	ej deeje	eejejf  f fddZ  ZS )Zamba2MambaDecoderLayerrZ   r   c                    s4   t    t||d| _t|j|jd| _|| _d S )N)rZ   r   r  )	r*   r+   r9  mambarU   r3   r  r  r   )r2   rZ   r   r5   r7   r8   r+     s   

z Zamba2MambaDecoderLayer.__init__r  r  r  r  NFrG   r  r   r   r  	use_cacher   transformer_hidden_statesr   c                 K   sd   |}|
dur
||
 n|}|  |}| j|||d}d}|| }|f}|r)||f7 }|r0||f7 }|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_values (`Zamba2HybridDynamicCache`, *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)rG   rZ  r   )r  r  )r2   rG   r  r   r   r   r  r  r  r   r  r   residualr  r  r7   r7   r8   rO     s"   


zZamba2MambaDecoderLayer.forward)	NNNNNFFNN)rQ   rR   rS   r   rl   r+   r   r-   r   r   rY   r2  r   rV   r  rO   rT   r7   r7   r5   r8   r    sF    	
r  c                       s   e Zd Zdedejdef fddZedddd		
	
	
	
	
			
dde	j
dee	j
 dee dee	j
 dee	j
 dee dee dee dee	j dee	jeee	je	jf  f fddZ  ZS )Zamba2HybridLayershared_transformerlinearr  c                    s    t    || _|| _|| _d S r)   )r*   r+   r  mamba_decoderr  )r2   r  r  r  r5   r7   r8   r+   Q  s   

zZamba2HybridLayer.__init__r  r  r  r  NFrG   r  r   r   r   r  r  r  r   c
              	   C   sn   | j |||||||	d}
|
d }|r|
d }| |}| j|||||||	d}
|r5|
d |f|
dd  }
|
S )aY  
        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_values (`Zamba2HybridDynamicCache`, *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`).
            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
                with `head_dim` being the embedding dimension of each attention head.
        )r  r   r   r  r  r  r   r   )r  r   r  r  r  r  r:   N)r  r  r  )r2   rG   r  r   r   r   r  r  r  r  layer_outputsr  r  r7   r7   r8   rO   Y  s2   !


zZamba2HybridLayer.forward)NNNNNFFN)rQ   rR   rS   r  r   r   r  r+   r   r-   r   r   rl   rY   r2  r   rV   r  rO   rT   r7   r7   r5   r8   r  P  sL    	
r  c                       sJ   e Zd ZU eed< dZdZddgZdZdZ	dZ
dZdZ fddZ  ZS )	Zamba2PreTrainedModelrZ   modelTr  r  r  c                    s   t  | t|tr_tt| jjt	
| jjt	
| jj  t	
| jj j| jjd}|t
t|   }|jj| td|jd }|jjt
| |jjd d S d S )N)minr   g      ?)r*   _init_weightsr   r9  r-   rj  randrZ   rs   mathrP  rH  rG  r   time_step_floorexpm1rN  datarl  rO  rE  rQ  rS  fill_)r2   r   rt  inv_dtrY  r5   r7   r8   r    s"   
z#Zamba2PreTrainedModel._init_weights)rQ   rR   rS   r   r   base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_flex_attn_supports_sdpa_is_statefulr  rT   r7   r7   r5   r8   r    s   
 r  c                       s   e Zd ZdZdef fddZe										ddeej	 deej
 deej	 d	ee d
eej dee dee dee dee deej	 deeef fddZdd Zdd Z  ZS )Zamba2Modelzh
    Model consisting of *config.num_hidden_layers* layers.

    Args:
        config: Zamba2Config
    rZ   c                    sN  t     | _ j| _ j| _t j j| j| _	 fddt
 jD }g }g } j| _t
 jD ]2} j| dkrH|t |d q5 j| dkrg|tj| jj| jjdd |t |d q5t|}t|}t|}| |||}t|| _ j| _t j jd| _ jr jrtd	 t | _d| _ | !  d S )
Nc                    s   g | ]}t  |d qS ))r   )r  )re   r   rZ   r7   r8   rh     s    z(Zamba2Model.__init__.<locals>.<listcomp>r  r   r^   Fr   r  ze`use_long_context` set to `True`: using rescaled `rope_theta` and extended `max_position_embeddings`.)"r*   r+   rZ   pad_token_idpadding_idx
vocab_sizer   	Embeddingr3   embed_tokensrz   r  rj   r{   r   r  r   iterr   
get_layersr  layersr  rU   r  final_layernormr  use_long_contextrV  rW  r   
rotary_embgradient_checkpointing	post_init)r2   rZ   blocksmamba_layerslinear_layersr   r  r5   r  r8   r+     s>   
zZamba2Model.__init__N	input_idsr   r   r  inputs_embedsr  r  output_hidden_statesreturn_dictr   r   c                 C   sd  |d ur|n| j j}|d ur|n| j j}|d ur|n| j j}|	d ur$|	n| j j}	|d u |d uA r4td| jrC| jrC|rCt	d d}|d u rL| 
|}|}t|}|rr|d u rr|d urb|jd n|jd }t| j || j| jd}|
d u r|d ur|j| jdnd}tj|||jd  |jd}
|d u r|
d}| |||
}| j jr| ||}nd }|rd	nd }|rd	nd }t| jD ]C\}}|r||f7 }| jr| jr| |j|||||||||
}n||||||||||d
	}|d }|r|d d ur||d f7 }q| |}|r||f7 }|d ur|jsd|_t||r#|nd ||d}|	r.|S | 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`.Fr   r=   r\   r  r   ra   r7   )r  r   r   r   r  r  r  r  T)last_hidden_stater  rG   
attentions) rZ   r  r  r  use_return_dict
ValueErrorr  r   rV  rW  r  r-   r  rB   rY   r=   r\   r   first_transformer_layer_idrO  r   _update_causal_maskr  r  r  r  _gradient_checkpointing_func__call__r  rk   r   to_tuple)r2   r  r   r   r  r  r  r  r  r  r   rG   r  r[   past_seen_tokensr   r  all_hidden_statesall_self_attnsr   layerr  r  r7   r7   r8   rO     s   





zZamba2Model.forwardc                 C   sv  | j jdkr|d urd|v r|S d S |j|j}}t|j}|jd }|d d }tj||f|||d}	|dkr@tj	|	dd}	|	tj
||d|ddk9 }	|	d d d d d d f |jd ddd}	|d ur|	 }	| d	kr|jd }
|	d
d |
f d|d d d d d d f d }|	d
d |
f |||	d
d |
f< | j jdkr|d ur|jjdv rt|	|}	|	S )Nflash_attention_2r   r   r;   )
fill_valuer=   r\   r.  ra   r   r:   .sdpa)r  xpunpu)rZ   r  r=   r\   r-   finfor  rB   fulltriurO  r   r   r  r   eqr3  r   r   _unmask_unattended)r2   r   r#  r   r=   r\   	min_dtypesequence_lengthtarget_lengthr   mask_lengthpadding_maskr7   r7   r8   r  ^  s0   
*
4$zZamba2Model._update_causal_maskc                 C   sp  g }g | _ d| _t| jD ]\}}|dkr| jdkr|| _t|}| jjt| jj dkrd| d}t	
|d d d d	 d
 }	| j |	 d}
| jD ]$}|dkrm|
| jj |jkrmt	
dt|
 d }| j | |
d7 }
qM| jjrd}
| jD ]$}|dkr|
| jj |jkrt	
dt|
 d }| j | |
d7 }
q{|t|t|t| q|t| q|S )Nr   r^   r   z	^layers\.z\.shared_transformer\.z(?:z3self_attn\.(?:q_proj|k_proj|v_proj|o_proj)\.weight|z1feed_forward\.(?:gate_up_proj|down_proj)\.weight|z,(?:input_layernorm|pre_ff_layernorm)\.weightz)$z>^shared_transformer\.feed_forward\.gate_up_proj_adapter_list\.z\.(?:0|1)\.weight$zg^shared_transformer\.self_attn\.(?:linear_q_adapter_list|linear_k_adapter_list|linear_v_adapter_list)\.)_tied_weights_keysr  r  rj   nextrZ   r  r   r  recompiler   r   r   r  r  )r2   r  r  r  r  layer_id
layer_typeblockprefix_patternmain_keys_pattern
adapter_id_layer_typeadapter_patternattn_adapter_patternr7   r7   r8   r    sh   




zZamba2Model.get_layers
NNNNNNNNNN)rQ   rR   rS   r   r   r+   r   r   r-   r   r   rY   r  r2  r   rV   r   rO   r  r  rT   r7   r7   r5   r8   r    sN    $	

x#r  c                       s   e Zd Zdef fddZe												ddeej deej	 deej d	ee
 d
eej deej dee dee dee dee deej deeej	f deeef fddZ						dddZ  ZS )Zamba2ForCausalLMrZ   c                    sP   t  | t|| _dg| jj| _|j| _tj|j|jdd| _	| 
  d S )Nzlm_head.weightFr   )r*   r+   r  r  r  r  r   r   r3   lm_headr  r2   rZ   r5   r7   r8   r+     s   
zZamba2ForCausalLM.__init__Nr   r  r   r   r  r  labelsr  r  r  r  r   logits_to_keepr   c                 K   s   |dur|n| j j}|	dur|	n| j j}	|
dur|
n| j j}
| j||||||||	||
d
}|d }t|tr<t| dn|}| |dd|ddf }d}|dur^| j	||| j
fi |}|
st|f|dd  }|durr|f| S |S t|||j|j|jdS )al  
        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]`.

        Example:

        ```python
        >>> from transformers import AutoTokenizer, Zamba2ForCausalLM

        >>> model = Zamba2ForCausalLM.from_pretrained("Zyphra/Zamba2-7B-v1")
        >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-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)
r  r   r   r  r  r  r  r  r   r  r   r   losslogitsr  rG   r  )rZ   r  r  r  r  r   rl   slicer#  loss_functionr  r   r  rG   r  )r2   r  r   r   r  r  r%  r  r  r  r  r   r&  r   r  rG   slice_indicesr)  r(  r  r7   r7   r8   rO     s@   (zZamba2ForCausalLM.forwardTc              	   K   s<  |d u }	|	s5|d us|d |j d kr"|d d |j d  d f }n!|j d |j d kr4|d d |f }nt| j|j d | j| jd}|d url|d u rl| dd }||dkd |	sl|d d |j d  d f }|d urw|	rwd|i}
nd| i}
|
	||||| jj
|d | D ]\}}||
vr||
|< q|
S )Nr;   r   r   r  r  r  )r   r  r  r   r&  r   )rB   rY   rZ   r=   r\   longr4  masked_fill_r   r   num_logits_to_keepitems)r2   r  r  r   r  r   r   r  r   empty_past_kvmodel_inputsr   r   r7   r7   r8   prepare_inputs_for_generation  sB   
z/Zamba2ForCausalLM.prepare_inputs_for_generation)NNNNNNNNNNNr   )NNNNNT)rQ   rR   rS   r   r+   r   r   r-   r   r   rY   r  r2  r   rl   rV   r   rO   r3  rT   r7   r7   r5   r8   r"    sb    
	

Tr"  a  
    The Zamba2 Model with a sequence classification head on top (linear layer).

    [`Zamba2ForSequenceClassification`] 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).
    )custom_introc                       s   e Zd Z fddZe										ddeej deej deej dee	e
eej f  deej d	eej d
ee dee dee dee de	eef fddZ  ZS )Zamba2ForSequenceClassificationc                    sJ   t  | |j| _t|| _| jj| _tj|j| jdd| _	| 
  d S )NFr   )r*   r+   
num_labelsr  r  r  r   r   r3   scorer  r$  r5   r7   r8   r+   a  s   

z(Zamba2ForSequenceClassification.__init__Nr  r   r   r  r  r%  r  r  r  r  r   c                 C   sB  |
dur|
n| j j}
| j||||||||	|
d	}|d }| |}|dur+|jd }n|jd }| j jdu r>|dkr>td| j jdu rGd}n1|durl|| j jk|jt	j
}t	j|jd |jt	j
d}|| d}nd}t| jj d |t	j||jd	|f }d}|dur||j}| j jdu r| jdkrd
| j _n| jdkr|jt	jks|jt	jkrd| j _nd| j _| j jd
krt }| jdkr|| | }n+|||}n%| j jdkrt }||d| j|d}n| j jdkrt }|||}|
s|f|dd  }|dur|f| S |S t|||j|j|jd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)r   r   r  r  r  r  r  r  r   r   z=Cannot handle batch sizes > 1 if no padding token is defined.r;   r]   z will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`ra   
regressionsingle_label_classificationmulti_label_classificationr'  )rZ   r  r  r7  rB   r  r  r>   r\   r-   int32rO  argmaxrV  rW  r6   rQ   problem_typer6  r=   r-  rl   r
   rh  r	   rC   r   r   r  rG   r  )r2   r  r   r   r  r  r%  r  r  r  r  transformer_outputsrG   r)  r[   last_non_pad_tokennon_pad_masktoken_indicespooled_logitsr(  loss_fctr  r7   r7   r8   rO   k  sx   



"


z'Zamba2ForSequenceClassification.forwardr!  )rQ   rR   rS   r+   r   r   r-   r   r   r   r   listr  r2  rV   r   rO   rT   r7   r7   r5   r8   r5  R  sH    
	

r5  )r"  r5  r  r  )r   )Nr   )Vr  r  	itertoolsr   typingr   r   r   r   r-   r   torch.nnr   r	   r
   activationsr   cache_utilsr   
generationr   modeling_attn_mask_utilsr   modeling_flash_attention_utilsr   modeling_outputsr   r   r   modeling_rope_utilsr   r   modeling_utilsr   r   processing_utilsr   utilsr   r   utils.deprecationr   utils.import_utilsr   r   configuration_zamba2r   +mamba_ssm.ops.triton.selective_state_updater    !mamba_ssm.ops.triton.ssd_combinedr!   r"   causal_conv1dr$   r%   
get_loggerrQ   rV  Moduler'   rU   rY   r   r   rl   r   r   r   r   r   r   r*  r,  r8  rk  rU  r9  r  r  r  r  r  r  r"  r5  __all__r7   r7   r7   r8   <module>   s   

m$

    /*AEK {  j