class LlamaRMSNorm(nn.Module):    def __init__(self, hidden_size, eps=1e-6):        """        LlamaRMSNorm is equivalent to T5LayerNorm        """        super().__init__()        self.weight = nn.Parameter(torch.ones(hidden_size))        self.variance_epsilon = eps # eps 防止取倒数之后分母为0    def forward(self, hidden_states):        input_dtype = hidden_states.dtype        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)        # weight 是末尾乘的可训练参数, 即g_i        return (self.weight * hidden_states).to(input_dtype)

class LlamaRotaryEmbedding(torch.nn.Module):class LlamaRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):        super().__init__()        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))        self.register_buffer("inv_freq", inv_freq)        # Build here to make `torch.jit.trace` work.        self.max_seq_len_cached = max_position_embeddings        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.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)        dtype = torch.get_default_dtype()        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)            def forward(self, x, seq_len=None):        # x: [bs, num_attention_heads, seq_len, head_size]        # This `if` block is unlikely to be run after we build sin/cos in `__init__`.        # Keep the logic here just in case.        if seq_len > self.max_seq_len_cached:            self.max_seq_len_cached = seq_len            t = torch.arange(self.max_seq_len_cached, device=x.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).to(x.device)            self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(x.dtype),            persistent=False)            self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(x.dtype),            persistent=False)            return (        self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),        self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),        )    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)
    def apply_rotary_pos_emb(q, k, cos, sin, position_ids):        # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.        cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]        sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]        cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]        sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]        q_embed = (q * cos) + (rotate_half(q) * sin)        k_embed = (k * cos) + (rotate_half(k) * sin)        return q_embed, k_embed

class LlamaDecoderLayer(nn.Module):    def __init__(self, config: LlamaConfig):    super().__init__()    self.hidden_size = config.hidden_size    self.self_attn = LlamaAttention(config=config)    self.mlp = LlamaMLP(        hidden_size=self.hidden_size,        intermediate_size=config.intermediate_size,        hidden_act=config.hidden_act,    )    self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)    self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    def forward(            self,            hidden_states: torch.Tensor,            attention_mask: Optional[torch.Tensor] = None,            position_ids: Optional[torch.LongTensor] = None,            past_key_value: Optional[Tuple[torch.Tensor]] = None,            output_attentions: Optional[bool] = False,            use_cache: Optional[bool] = False,    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:        residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)    # Self Attention    hidden_states, self_attn_weights, present_key_value = self.self_attn(        hidden_states=hidden_states,        attention_mask=attention_mask,        position_ids=position_ids,        past_key_value=past_key_value,        output_attentions=output_attentions,        use_cache=use_cache,    )    hidden_states = residual + hidden_states    # Fully Connected    residual = hidden_states    hidden_states = self.post_attention_layernorm(hidden_states)    hidden_states = self.mlp(hidden_states)    hidden_states = residual + hidden_states    outputs = (hidden_states,)    if output_attentions:        outputs += (self_attn_weights,)    if use_cache:        outputs += (present_key_value,)    return outputs

class MultiQueryAttention(nn.Module):"""Multi-Query self attention.Using torch or triton attention implemetation enables user to also useadditive bias."""    def __init__(        self,        d_model: int,        n_heads: int,        device: Optional[str] = None,    ):        super().__init__()        self.d_model = d_model        self.n_heads = n_heads        self.head_dim = d_model // n_heads        self.Wqkv = nn.Linear( # Multi-Query Attention 创建            d_model,            d_model + 2 * self.head_dim, # 只创建查询的头向量，所以只有1 个d_model            device=device, # 而键和值则共享各自的一个head_dim 的向量        )        self.attn_fn = scaled_multihead_dot_product_attention        self.out_proj = nn.Linear(            self.d_model,            self.d_model,            device=device        )        self.out_proj._is_residual = True # type: ignore    def forward(        self,        x,    ):        qkv = self.Wqkv(x) # (1, 512, 960)        query, key, value = qkv.split( # query -> (1, 512, 768)            [self.d_model, self.head_dim, self.head_dim], # key -> (1, 512, 96)            dim=2 # value -> (1, 512, 96)        )        context, attn_weights, past_key_value = self.attn_fn(            query,            key,            value,            self.n_heads,            multiquery=True,    )        return self.out_proj(context), attn_weights, past_key_value

# Multi Head Attentionself.Wqkv = nn.Linear( # Multi-Head Attention 的创建方法    self.d_model,    3 * self.d_model, # 查询、键和值3 个矩阵, 所以是3 * d_model    device=device)query, key, value = qkv.chunk( # 每个tensor 都是(1, 512, 768)    3,    dim=2)# Multi Query Attentionself.Wqkv = nn.Linear( # Multi-Query Attention 的创建方法    d_model,    d_model + 2 * self.head_dim, # 只创建查询的头向量，所以是1* d_model    device=device, # 而键和值不再具备单独的头向量)query, key, value = qkv.split( # query -> (1, 512, 768)    [self.d_model, self.head_dim, self.head_dim], # key -> (1, 512, 96)    dim=2 # value -> (1, 512, 96))