dearth_model.py

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch import Tensor
from typing import Optional, Tuple
import math

import logging
import copy

from dearth_config import DearthConfig

_USE_FAST_ROPE = False

class RMSNorm(torch.nn.Module): # a variant of LayerNorm that is faster and more memory efficient
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        # set the weight to be 1 initially
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding
class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )

        self.register_buffer("default_pos_ids", 
                             torch.arange(0, self.max_position_embeddings, dtype=torch.long).view(-1, self.max_position_embeddings), 
                             persistent=False)

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) # shape: (max_seq_len_cached, dim // 2)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

# Copied from transformers.models.llama.modeling_llama.rotate_half
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)
# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    cos = cos[position_ids].unsqueeze(1)  # [seq_len, dim] -> [batch_size, 1, seq_len, head_dim]
    sin = sin[position_ids].unsqueeze(1)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed




class FastRope(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        cis = precompute_freqs_cis(dim, max_position_embeddings, theta=base)
        self.register_buffer("cis", cis, persistent=False)

    def forward(self, start_idx, seq_len):
        return self.cis[start_idx:start_idx+seq_len, :]

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """
    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.

    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'
    and the end index 'end'. The 'theta' parameter scales the frequencies.
    The returned tensor contains complex values in complex64 data type.

    Args:
        dim (int): Dimension of the frequency tensor.
        end (int): End index for precomputing frequencies.
        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.

    Returns:
        torch.Tensor: Precomputed frequency tensor with complex exponentials.
    """
    with torch.no_grad():
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
        t = torch.arange(end, device=freqs.device)  # type: ignore
        freqs = torch.outer(t, freqs).float()  # type: ignore
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
        return freqs_cis      

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Apply rotary embeddings to input tensors using the given frequency tensor.

    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
    returned as real tensors.

    Args:
        xq (torch.Tensor): Query tensor to apply rotary embeddings.
        xk (torch.Tensor): Key tensor to apply rotary embeddings.
        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
    """
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    """
    Reshape frequency tensor for broadcasting it with another tensor.

    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
    for the purpose of broadcasting the frequency tensor during element-wise operations.

    Args:
        freqs_cis (torch.Tensor): Frequency tensor to be reshaped.
        x (torch.Tensor): Target tensor for broadcasting compatibility.

    Returns:
        torch.Tensor: Reshaped frequency tensor.

    Raises:
        AssertionError: If the frequency tensor doesn't match the expected shape.
        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
    """
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1]), f"freqs_cis.shape: {freqs_cis.shape}, x.shape: {x.shape}"
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)




class AttentionMask(nn.Module):
    attn_mask: torch.Tensor = None
    def __init__(self, config: DearthConfig):
        super().__init__()
        self.config = config
        self.sliding_window_size = config.sliding_window_size
        self.front_window_size = config.front_window_size
        if self.attn_mask is None: 
            tmp_attn_mask = self.build_causal_and_window_mask(config.max_token_len, config.sliding_window_size, config.front_window_size)
            self.attn_mask = tmp_attn_mask.requires_grad_(False) # shape: (max_token_len, max_token_len)
        #self.register_buffer("attn_mask", self.build_causal_and_window_mask(config.max_token_len, config.sliding_window_size, config.front_window_size).requires_grad_(False), persistent=False)

    def forward(self, bz, n_head, q_seq_len, kv_seq_len, q_start_idx: int, device, dtype) -> torch.Tensor:
        if self.attn_mask.device != device or self.attn_mask.dtype != dtype:
            self.attn_mask = self.attn_mask.to(device=device, dtype=dtype).requires_grad_(False)
        end_idx = q_start_idx + q_seq_len
        q_k_diff_len = kv_seq_len - q_seq_len # it should be >= 0, because it is meaningless to attend future tokens
        top = q_start_idx
        bottom = end_idx
        if q_start_idx == 0 and q_k_diff_len == 0:
            # assume: sliding window size = 100, front window size = 50
            # case 1: training: q_start_idx = 0, q_seq_len = 1000, kv_seq_len = 1000
            mask = self.attn_mask[:end_idx, :end_idx]
        elif q_k_diff_len > 0 and q_start_idx > 0 and end_idx >= kv_seq_len:
            # TODO: not allow in training; remove this line after testing
            raise RuntimeError(f"NOT FOR TRAINING: q_start_idx = {q_start_idx}, q_seq_len = {q_seq_len}, kv_seq_len = {kv_seq_len}")
            if end_idx > self.front_window_size + self.sliding_window_size:
                # case 2: qsl < kvsl: q_start_idx = 190, q_seq_len = 10, kv_seq_len = 150, end_idx = 200
                # mask = self.attn_mask[top:bottom, :self.front_window_size] + \
                #     self.attn_mask[q_start_idx:end_idx, end_idx - (kv_seq_len - self.front_window_size):end_idx]
                mask = torch.cat([self.attn_mask[top:bottom, :self.front_window_size], self.attn_mask[top:bottom, end_idx - (kv_seq_len - self.front_window_size):end_idx]], dim=-1)
            elif end_idx <= self.front_window_size + self.sliding_window_size:
                # case 3: qsl < kvsl: q_start_idx = 140, q_seq_len = 10, kv_seq_len = 150, end_idx = 150
                mask = self.attn_mask[top:bottom, :end_idx]
        else:
            raise RuntimeError(f"q_start_idx = {q_start_idx}, q_seq_len = {q_seq_len}, kv_seq_len = {kv_seq_len}")
        return mask.expand(bz, n_head, q_seq_len, kv_seq_len).detach()
        
    
    @staticmethod
    def build_causal_and_window_mask(seq_len, sliding_window_size, front_window_size) -> torch.Tensor:
        mask = torch.ones(seq_len, seq_len)
        if seq_len > sliding_window_size: # need to apply sliding window mask, beacause the sequence is too long
            mask = torch.triu(mask, diagonal=-sliding_window_size+1)
            if front_window_size > 0:
                tmp_front_mask = torch.cat([torch.ones(seq_len, front_window_size), torch.zeros(seq_len, seq_len-front_window_size)], dim=-1)
                tmp_front_mask = torch.tril(tmp_front_mask, diagonal=-sliding_window_size)
                mask = mask + tmp_front_mask
        # apply causal mask
        mask = mask.tril(diagonal=0)
        mask = mask.log() # map 0 to -inf, 1 to 0
        # print(f"mask.shape: {mask.shape}, and mask")
        # print(mask)
        return mask
    

class SharedAttentionMask(nn.Module):
    def __init__(self, config: DearthConfig):
        super().__init__()
        self.config = config
        self.sliding_window_size = config.sliding_window_size
        self.front_window_size = config.front_window_size
        tmp_attn_mask = self.build_causal_and_window_mask(config.max_token_len, config.sliding_window_size, config.front_window_size)
        self.register_buffer("attn_mask", tmp_attn_mask, persistent=False)

    def forward(self, q_seq_len, kv_seq_len, q_start_idx: int) -> torch.Tensor:
        end_idx = q_start_idx + q_seq_len
        q_k_diff_len = kv_seq_len - q_seq_len # it should be >= 0, because it is meaningless to attend future tokens
        top = q_start_idx
        bottom = end_idx
        if q_start_idx == 0 and q_k_diff_len == 0:
            # assume: sliding window size = 100, front window size = 50
            # case 1: training: q_start_idx = 0, q_seq_len = 1000, kv_seq_len = 1000
            mask = self.attn_mask[:end_idx, :end_idx]
        elif q_k_diff_len > 0 and q_start_idx > 0 and end_idx >= kv_seq_len:
            # TODO: not allow in training; remove this line after testing
            raise RuntimeError(f"NOT FOR TRAINING: q_start_idx = {q_start_idx}, q_seq_len = {q_seq_len}, kv_seq_len = {kv_seq_len}")
            if end_idx > self.front_window_size + self.sliding_window_size:
                # case 2: qsl < kvsl: q_start_idx = 190, q_seq_len = 10, kv_seq_len = 150, end_idx = 200
                # mask = self.attn_mask[top:bottom, :self.front_window_size] + \
                #     self.attn_mask[q_start_idx:end_idx, end_idx - (kv_seq_len - self.front_window_size):end_idx]
                mask = torch.cat([self.attn_mask[top:bottom, :self.front_window_size], self.attn_mask[top:bottom, end_idx - (kv_seq_len - self.front_window_size):end_idx]], dim=-1)
            elif end_idx <= self.front_window_size + self.sliding_window_size:
                # case 3: qsl < kvsl: q_start_idx = 140, q_seq_len = 10, kv_seq_len = 150, end_idx = 150
                mask = self.attn_mask[top:bottom, :end_idx]
        else:
            raise RuntimeError(f"q_start_idx = {q_start_idx}, q_seq_len = {q_seq_len}, kv_seq_len = {kv_seq_len}")
        return mask.detach() # shape: (1, 1, seqlen, seqlen)
        
    
    @staticmethod
    def build_causal_and_window_mask(seq_len, sliding_window_size, front_window_size) -> torch.Tensor:
        mask = torch.ones(seq_len, seq_len)
        if seq_len > sliding_window_size: # need to apply sliding window mask, beacause the sequence is too long
            mask = torch.triu(mask, diagonal=-sliding_window_size+1)
            if front_window_size > 0:
                tmp_front_mask = torch.cat([torch.ones(seq_len, front_window_size), torch.zeros(seq_len, seq_len-front_window_size)], dim=-1)
                tmp_front_mask = torch.tril(tmp_front_mask, diagonal=-sliding_window_size)
                mask = mask + tmp_front_mask
        # apply causal mask
        mask = mask.tril(diagonal=0)
        mask = mask.log() # map 0 to -inf, 1 to 0
        # print(f"mask.shape: {mask.shape}, and mask")
        # print(mask)
        return mask



def build_mpt_alibi_tensor(num_heads, sequence_length, alibi_bias_max=8, device=None):
    r"""
    Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
    relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
    the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
    https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292

    retrun shape: (1, num_heads, 1, sequence_length)
    """
    alibi = torch.arange(1 - sequence_length, 1, dtype=torch.int32, device=device).view(1, 1, 1, sequence_length)
    num_heads_power_of_2 = 2 ** math.ceil(math.log2(num_heads))

    base = torch.arange(1, num_heads_power_of_2 + 1, dtype=torch.float32, device=device)
    base = base * (alibi_bias_max / num_heads_power_of_2)

    slopes = 1.0 / torch.pow(2, base)
    slopes = slopes.view(1, num_heads, 1, 1)

    if num_heads_power_of_2 != num_heads:
        slopes = torch.concat([slopes[1::2], slopes[::2]])[:num_heads]

    alibi = alibi * slopes
    return alibi


# def build_alibi_tensor(num_heads, sequence_length, alibi_bias_max=8, device=None):
#     r"""
#     Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
#     relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
#     the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
#     https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292

#     retrun shape: (1, num_heads, 1, sequence_length)
#     """
#     slope = []
#     m_power = (-8/num_heads)
#     m_increace = -8/num_heads
#     for i in range(num_heads):
#         slope.append(m_power)
#         m_power += m_increace
#     slope = torch.tensor(slope, device=device)
#     alibi = torch.arange(1 - sequence_length, 1, dtype=torch.int32, device=device).view(1, 1, 1, sequence_length)
#     alibi = alibi * slope.view(1, num_heads, 1, 1)
#     return alibi

def compute_alibi(num_heads, sequence_length, alibi_bias_max=8, device=None):
    r"""
    Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
    relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
    the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
    https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292

    retrun shape: (1, num_heads, 1, sequence_length)
    """
    slope = []
    m_power = (-8/num_heads)
    m_increace = -8/num_heads
    for i in range(num_heads):
        slope.append(2 ** m_power)
        m_power += m_increace
    slope = torch.tensor(slope, device=device)
    alibi = torch.arange(1 - sequence_length, 1, dtype=torch.int32, device=device).view(1, 1, 1, sequence_length)
    alibi = alibi * slope.view(1, num_heads, 1, 1)
    return alibi


class Attention(nn.Module):
    def __init__(self, config: DearthConfig):
        super().__init__()
        assert config.dim % config.n_head == 0

        # regularization
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head if config.n_kv_head is not None else config.n_head
        self.dim = config.dim
        assert config.dim % config.n_head == 0
        self.dim_qk_head = config.dim_qk_head if config.dim_qk_head is not None else config.dim // config.n_head
        self.dim_v_head = config.dim // config.n_head
        assert config.n_kv_head <= config.n_head and config.n_head % config.n_kv_head == 0
        self.n_kv_group = config.n_head // config.n_kv_head
        self.dropout_rate = config.dropout_rate

        self.alibi_emb = None
        self.pos_emb = None

        self.sliding_window_size = config.sliding_window_size

        def _fill_with_neg_inf(t):
            """FP16-compatible function that fills a tensor with -inf."""
            return t.float().fill_(float("-inf")).type_as(t)
        
        # neg_inf_mask = _fill_with_neg_inf(torch.ones_like(torch.empty(config.max_token_len, config.max_token_len)))
        # window_size_mask = torch.triu(neg_inf_mask, diagonal=1)
        # if config.sliding_window_size is not None and config.max_token_len > config.sliding_window_size:
        #     window_size_mask = window_size_mask + torch.tril(neg_inf_mask, diagonal=-config.sliding_window_size)
        # self.register_buffer("window_size_mask", window_size_mask, persistent=False)
        # if config.use_alibi:
        #     alibi_emb = compute_alibi(config.n_head, config.max_token_len) # shape: (1, n_head, 1, seqlen)
        #     #self.alibi_emb = self.alibi_emb.expand(1, config.n_head, config.max_token_len, config.max_token_len) # shape: (1, n_head, seqlen, seqlen)
        #     self.register_buffer("alibi_emb", alibi_emb, persistent=False)

        self.window_size_mask = AttentionMask(config)

        if config.use_rotary:
            if not _USE_FAST_ROPE:
                self.pos_emb = RotaryEmbedding(
                    self.dim_qk_head,
                    max_position_embeddings=config.max_token_len,
                    base=config.rope_theta,
                )
            if _USE_FAST_ROPE:
                self.pos_emb = FastRope(
                    self.dim_qk_head,
                    max_position_embeddings=config.max_token_len,
                    base=config.rope_theta,
                )

        # query, key, values projections for all heads
        self.wq = nn.Linear(self.dim, self.n_head * self.dim_qk_head, bias=True)
        self.wk = nn.Linear(self.dim, self.n_kv_head * self.dim_qk_head, bias=True)
        self.wv = nn.Linear(self.dim, self.dim // self.n_kv_group, bias=False)
        self.wo = nn.Linear(self.dim, self.dim, bias=False)

        
    def forward(self, x: Tensor, attn_mask: Tensor, start_idx: Optional[int] = 0):
        batch_size, seqlen, emb_dim = x.size() # batch size, sequence length, embedding dimensionality (dim)
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        # split embedding dim into number of heads
        xq = xq.view(batch_size, seqlen, self.n_head, self.dim_qk_head)
        xk = xk.view(batch_size, seqlen, self.n_kv_head, self.dim_qk_head)
        xv = xv.view(batch_size, seqlen, self.n_kv_head, self.dim_v_head)

        if self.pos_emb is not None and _USE_FAST_ROPE:
            xq, xk = apply_rotary_emb(xq, xk, self.pos_emb(start_idx, seqlen))

        # transpose to get dimensions batch_size * n_head * seqlen * emb_dim
        xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
        kv_seqlen = xk.size(2)

        # apply positional embeddings
        if self.pos_emb is not None and not _USE_FAST_ROPE:
            # self.pos_emb = self.pos_emb.to(x.device, dtype=x.dtype)
            # xq, xk = apply_rotary_pos_emb(xq, xk, self.pos_emb[start_idx:start_idx+seqlen])
            cos, sin = self.pos_emb(xv, seq_len=kv_seqlen)
            xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin, self.pos_emb.default_pos_ids[:, :kv_seqlen])

        # TODO: add cache for fast inference


        # grouped query
        xk = repeat_kv(xk, self.n_kv_group)
        xv = repeat_kv(xv, self.n_kv_group)

        # self.window_size_mask = self.window_size_mask.to(x.device, dtype=x.dtype)
        # attn_mask = self.window_size_mask[start_idx:start_idx+seqlen, start_idx:start_idx+kv_seqlen]
        # attn_mask = attn_mask.unsqueeze(0).unsqueeze(0) # shape: (1, 1, seqlen, seqlen)
        # attn_mask = attn_mask.expand(batch_size, self.n_head, seqlen, kv_seqlen) # shape: (batch_size, n_head, seqlen, seqlen)
        # if self.alibi_emb is not None:
        #     self.alibi_emb = self.alibi_emb.to(x.device, dtype=x.dtype)
        #     attn_mask = attn_mask + self.alibi_emb[:,:,:,:kv_seqlen]

        #attn_mask = self.window_size_mask(batch_size, self.n_head, seqlen, kv_seqlen, start_idx, x.device, x.dtype) # -inf or 0
        
        # efficient attention using Flash Attention CUDA kernels
        y = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=self.dropout_rate if self.training else 0)
        y = y.transpose(1, 2).contiguous().view(batch_size, seqlen, emb_dim) # merge heads

        # output projection
        return self.wo(y)

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    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)
    hidden_states.shape = (batch, n_kv_head, seqlen, head_dim)
    """
    # if n_rep == 1:
    #     return hidden_states
    # return torch.repeat_interleave(hidden_states, n_rep, dim=1)

    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

# def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
#     """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
#     bs, slen, n_kv_heads, head_dim = x.shape
#     if n_rep == 1:
#         return x
#     return (
#         x[:, :, :, None, :]
#         .expand(bs, slen, n_kv_heads, n_rep, head_dim)
#         .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
#     )

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        dim = config.dim
        hidden_dim = config.dim * 4 if config.hidden_dim is None else config.hidden_dim
        multiple_of = 64 if config.multiple_of is None else config.multiple_of
        hidden_dim = int(2 * hidden_dim / 3)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) # round up to nearest multiple of multiple_of

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Mimic_Attn(Attention):
    def __init__(self, config):
        new_config = copy.deepcopy(config)
        new_config.n_head = config.mimic_n_head if config.mimic_n_head is not None else config.n_head
        new_config.n_kv_head = config.mimic_n_kv_head if config.mimic_n_kv_head is not None else config.n_kv_head
        new_config.dim_qk_head = config.mimic_dim_qk_head if config.mimic_dim_qk_head is not None else config.dim_qk_head
        new_config.dropout_rate = config.mimic_attn_dropout if config.mimic_attn_dropout is not None else 0.0
        new_config.use_rotary = config.mimic_use_rotary if config.mimic_use_rotary is not None else config.use_rotary
        new_config.use_alibi = config.mimic_use_alibi if config.mimic_use_alibi is not None else config.use_alibi

        super().__init__(new_config)
        self.saved_q = None
        self.saved_k = None
        self.saved_v = None
        self.saved_attn_map = None

    def forward(self, x: Tensor, attn_mask: Tensor, start_idx: Optional[int] = 0): # shape of attn_mask: (bz, n_head, q_seq_len, kv_seq_len)
        batch_size, seqlen, emb_dim = x.size() # batch size, sequence length, embedding dimensionality (dim)
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        self.saved_v = xv

        # split embedding dim into number of heads
        xq = xq.view(batch_size, seqlen, self.n_head, self.dim_qk_head)
        xk = xk.view(batch_size, seqlen, self.n_kv_head, self.dim_qk_head)
        xv = xv.view(batch_size, seqlen, self.n_kv_head, self.dim_v_head)

        if self.pos_emb is not None and _USE_FAST_ROPE:
            xq, xk = apply_rotary_emb(xq, xk, self.pos_emb(start_idx, seqlen))

        # transpose to get dimensions batch_size * n_head * seqlen * emb_dim
        xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
        kv_seqlen = xk.size(2)

        # # apply positional embeddings
        # if self.pos_emb is not None:
        #     self.pos_emb = self.pos_emb.to(x.device)
        #     xq, xk = apply_pos_emb(xq, xk, self.pos_emb[start_idx:start_idx+seqlen])
        if self.pos_emb is not None and not _USE_FAST_ROPE:
            cos, sin = self.pos_emb(xv, seq_len=kv_seqlen)
            xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin, self.pos_emb.default_pos_ids[:, :kv_seqlen])

        # TODO: add cache for fast inference

        # grouped query
        xk = repeat_kv(xk, self.n_kv_group)
        xv = repeat_kv(xv, self.n_kv_group)

        # self.window_size_mask = self.window_size_mask.to(x.device)
        # kv_seqlen = xk.size(2)
        # attn_mask = self.window_size_mask[start_idx:start_idx+seqlen, start_idx:start_idx+kv_seqlen]
        # attn_mask = attn_mask.unsqueeze(0).unsqueeze(0) # shape: (1, 1, seqlen, seqlen)
        # attn_mask = attn_mask.expand(batch_size, self.n_head, seqlen, kv_seqlen) # shape: (batch_size, n_head, seqlen, seqlen)
        # if self.alibi_emb is not None:
        #     self.alibi_emb = self.alibi_emb.to(x.device)
        #     attn_mask = attn_mask + self.alibi_emb[:,:,:,:kv_seqlen]

        #attn_mask = self.window_size_mask(batch_size, self.n_head, seqlen, kv_seqlen, start_idx, x.device, x.dtype) # -inf or 0

        attn_weights = torch.matmul(xq, xk.transpose(2, 3)) * (1 / math.sqrt(self.dim_qk_head)) # shape: (batch_size, n_head, seqlen, seqlen)
        attn_weights = attn_weights + attn_mask.expand(batch_size, self.n_head, seqlen, kv_seqlen) # shape: (batch_size, n_head, seqlen, seqlen
        attn_weights = F.softmax(attn_weights.float(), dim=-1).to(xq.dtype) # shape: (batch_size, n_head, seqlen, seqlen)
        # use log_softmax to avoid overflow
        #attn_weights = F.log_softmax(attn_weights, dim=-1).exp() # shape: (batch_size, n_head, seqlen, seqlen)
        self.saved_attn_map = attn_weights

        attn_weights = F.dropout(attn_weights, p=self.dropout_rate, training=self.training)

        y = torch.matmul(attn_weights, xv) # shape: (batch_size, n_head, seqlen, head_dim)

        y = y.transpose(1, 2).contiguous().view(batch_size, seqlen, emb_dim) # merge heads

        # output projection
        return self.wo(y)

    def get_intermediate_attn_v(self):
        return self.saved_attn_map, self.saved_v


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = RMSNorm(config.dim)
        self.attn = Attention(config)
        self.ln_2 = RMSNorm(config.dim)
        self.mlp = MLP(config)

        self.residual_factor = config.residual_factor

    def forward(self, x: Tensor, attn_mask: Tensor, start_idx: int):
        # post-LN
        residual = x
        x = self.attn(x, attn_mask, start_idx=start_idx)
        x = self.ln_1(self.residual_connection(x, residual))

        residual = x
        x = self.mlp(x)
        x = self.ln_2(self.residual_connection(x, residual))

        return x
    
    def residual_connection(self, x, residual):
        # residual factor should > 1.0
        return residual * self.residual_factor + x



class DearthModel(nn.Module):
    def __init__(self, config: DearthConfig):
        super().__init__()
        assert config.vocab_size is not None
        assert config.max_token_len is not None

        self.layer_init_factor = config.layer_init_factor if config.layer_init_factor is not None else float(config.n_layer * 8) ** (-1/2)
        self.residual_factor = config.residual_factor if config.residual_factor is not None else float(config.n_layer * 2) ** (1/4)
        if config.residual_factor is None:
            config.residual_factor = self.residual_factor
            #logging.warning(f"residual_factor is not set, using default value {self.residual_factor} = (2 * n_layer) ** 1/4")
        if config.layer_init_factor is None:
            config.layer_init_factor = self.layer_init_factor
            #logging.warning(f"layer_init_factor is not set, using default value {self.layer_init_factor} = (n_layer * 8) ** -1/2")
        
        self.config = config

        layers = []
        for i in range(config.n_layer):
            if config.mimic_attn_layer is not None and i+1 == config.mimic_attn_layer:
                new_layer = TransformerBlock(config)
                new_layer.attn = Mimic_Attn(config)
                layers.append(new_layer)
            else:
                layers.append(TransformerBlock(config))

        self.layers = nn.ModuleList(layers)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)
        self.ln_before = RMSNorm(config.dim)
        self.shared_attn_mask = SharedAttentionMask(config)

        if config.mimic_attn_layer is not None and config.mimic_attn_layer > 0 and config.mimic_attn_layer <= config.n_layer:
            self.mimic_attn = self.layers[config.mimic_attn_layer-1].attn
        else:
            self.mimic_attn = None

        # initialize weights
        _init_weight(self, self.layer_init_factor)

    def get_input_device(self):
        return self.embed_tokens.weight.device

    # def _init_weights(self, module):
    #     if isinstance(module, nn.Linear):
    #         torch.nn.init.xavier_normal_(module.weight, gain=self.layer_init_factor)
    #         if module.bias is not None:
    #             torch.nn.init.zeros_(module.bias)
    #     elif isinstance(module, nn.Embedding):
    #         torch.nn.init.xavier_normal_(module.weight, gain=1)
    #     elif isinstance(module, RMSNorm):
    #         module.weight.data.fill_(1.0)
            


    def forward(self, tokens, start_idx=0): # return all logits
        batch_size, seqlen = tokens.size()
        if seqlen > self.config.max_token_len:
            raise ValueError(f"input sequence length {seqlen} exceeds maximum sequence length {self.config.max_token_len}")

        # create token embeddings from token table; x.shape = (batch_size, seqlen, dim)
        h = self.embed_tokens(tokens)
        assert h.size() == (batch_size, seqlen, self.config.dim)

        h = self.ln_before(h)

        # transformer layers
        attn_mask = self.shared_attn_mask(seqlen, seqlen, q_start_idx=start_idx) # TODO: it will not work if q_seq_len != kv_seq_len
        for layer in self.layers:
            h = layer(h, attn_mask, start_idx=start_idx) # h.shape = (batch_size, seqlen, dim)

        return h, None


    def get_num_params(self):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        #n_params = sum(p.numel() for p in self.parameters())
        n_params = sum(p.numel() for p in self.transformer.layers[0].parameters() if p.requires_grad)
        return int(n_params)
    
    
    def get_intermediate_attn_v(self):
        if self.mimic_attn is None:
            return torch.zeros(1, 1, 1, 1), torch.zeros(1, 1, 1, 1)
        return self.mimic_attn.get_intermediate_attn_v()


class DearthForCausalLM(nn.Module):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: DearthConfig):
        super().__init__()
        self.model = DearthModel(config)
        self.dearth_config = config
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        torch.nn.init.xavier_normal_(self.lm_head.weight, gain=1)

        self.front_window_size = config.front_window_size
        self.sliding_window_size = config.sliding_window_size

    def get_input_device(self):
        return self.model.get_input_device()
    
    def get_intermediate_attn_v(self):
        return self.model.get_intermediate_attn_v()
    
    def print_all_params(self):
        for name, param in self.named_parameters():
            print(f"name: {name}, param.shape: {param.shape}")

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        use_cache: Optional[bool] = False,
    ) ->Tuple: #-> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        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]`.

        Returns:

        Example:

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

        >>> model = MistralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)

        >>> 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."
        ```"""
        outputs = self.model(
            tokens=input_ids
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        output = (logits,) + outputs[1:]
        return output


def _init_weight(model, weight_init_factor): # TODO: fix this part if change any model structure
    small_list = {'wv', 'wo', 'w1', 'w2', 'w3'}
    norm_list = {'ln_before', 'ln_2', 'ln_1'}
    for name, p in model.named_parameters():
        percise_name = name.split(".")[-2]
        if "bias" in name:
            logging.debug(f"the parameter {name} is initialized with 0.0")
            p.data.fill_(0.0)
        elif percise_name in small_list:
            logging.debug(f"the parameter {name} is initialized with gain={weight_init_factor}")
            torch.nn.init.xavier_normal_(p, gain=weight_init_factor)
        elif percise_name in norm_list:
            logging.debug(f"the parameter {name} is initialized with 1.0")
            p.data.fill_(1.0)
        else:
            logging.debug(f"the parameter {name} is initialized with gain=1.0")
            torch.nn.init.xavier_normal_(p, gain=1)