Tiny Universe - Llama3架构

Llama3和Llama2和Qwen2的整体架构相似，本篇文章主要讲解它们的一些主要不同点。

关于Qwen2架构可参考 Qwen2架构学习笔记

llama3区别于llama2在模型层面的区别主要体现在全模型使用GQA。

基础知识

MLP

MLP（Multi-Layer Perceptron）多层感知机是一种前馈神经网络，由一个或多个全连接层组成。每个全连接层包含一组可学习的权重矩阵和偏置向量，用于将输入数据进行线性变换和非线性激活。MLP可以用于各种任务，如分类、回归等。

在大模型中，MLP通常作为基本的网络组件，用于构建更复杂的结构。例如，在Transformer中，前馈神经网络部分就是一个MLP。此外，MLP还可以与其他网络结构（如卷积神经网络）结合，形成更强大的模型。

典型的MLP包括包括三层：输入层、隐层和输出层，MLP神经网络不同层之间是全连接的( 全连接的意思就是：上一层的任何一个神经元与下一层的所有神经元都有连接)。

如图所示

Attention

在自注意力机制中，输入序列的每个元素首先通过三个不同的线性变换，分别生成 Query（查询）、Key（键）、和 Value（值）矩阵。这三个矩阵共同用于计算输入序列中各个元素之间的注意力权重。

假设输入序列为 X=[x1,x2,…,xn]X = [x_1, x_2, \dots, x_n]X=[x1,x2,…,xn]，这些元素经过线性变换后得到：

Q=XWQ,K=XWK,V=XWVQ = XW^Q, \quad K = XW^K, \quad V = XW^VQ=XWQ,K=XWK,V=XWV

其中：

XXX 是输入序列，每个元素是一个向量。
WQ,WK,WVW^Q, W^K, W^VWQ,WK,WV 分别是用于生成 Query、Key、Value 的可学习权重矩阵。

1.1 点积注意力（Scaled Dot-Product Attention）

通过 Q、K 矩阵，计算每个输入与其他输入的相关性。具体公式如下：

$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^\top}{\sqrt{d_k}} \right) V$

其中：

Q 是查询矩阵。
K 是键矩阵， $K^\top$ 是键矩阵的转置。
V 是值矩阵。
$d_k$ 是键向量的维度， $\sqrt{d_k}$ 是缩放因子，用于避免点积值过大导致 softmax 输出过小的梯度。

1.2 计算步骤解析

生成 Query、Key、Value 矩阵：输入序列经过不同的线性变换，生成对应的 Q、K、V 矩阵。
计算点积：对 Query 和 Key 矩阵进行点积，得到输入序列中每个元素与其他元素之间的相关性分数。
缩放与归一化：将相关性分数除以 dk\sqrt{d_k}dk 进行缩放，并通过 softmax 归一化，得到注意力权重。
加权求和：将注意力权重与 Value 矩阵相乘，得到最终的上下文向量。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super(ScaledDotProductAttention, self).__init__()
        self.d_k = d_k

    def forward(self, Q, K, V, mask=None):
        # Q, K, V: batch_size x seq_len x d_k
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)  # Apply the mask (optional)
        attn_weights = F.softmax(scores, dim=-1)  # softmax over the last dimension
        output = torch.matmul(attn_weights, V)  # Weighted sum of values
        return output, attn_weights

多头注意力机制通过并行计算多个注意力头来捕捉不同子空间的特征。每个头独立生成自己的 Q、K、V 矩阵，进行自注意力计算，然后将各个头的结果拼接起来，通过一个线性层投影到最终输出。

2.1 多头注意力公式

对于多头注意力机制，公式如下：

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$

其中每个头的计算方式为：

$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

$W_i^Q, W_i^K, W_i^V$ 是每个头的可学习权重。
$W^O$ 是用于拼接后投影的线性变换矩阵。

多头注意力机制允许模型在不同的子空间中关注输入序列的不同部分，从而增强了模型的表达能力。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_k = d_model // num_heads
        self.num_heads = num_heads

        # Define weight matrices for Q, K, V and output projection
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.fc_out = nn.Linear(d_model, d_model)

        # Attention module
        self.attention = ScaledDotProductAttention(self.d_k)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # Perform linear transformation and split into multiple heads
        Q = self.W_Q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # Apply scaled dot-product attention to each head
        attn_output, attn_weights = self.attention(Q, K, V, mask)

        # Concatenate heads and apply final linear projection
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        output = self.fc_out(attn_output)

        return output, attn_weights

RMS正则化

RMSProp（Root Mean Square Propagation）是对学习率进行自适应调整的优化器，用来应对不稳定的梯度更新问题。它通过引入梯度平方的指数加权移动平均来动态调整每个参数的学习率。

1.1 RMSProp 的更新公式

对于每个参数 $\theta_i$ ，RMSProp 的更新公式如下：

$E[g^2]_t = \beta E[g^2]_{t-1} + (1 - \beta) g_t^2$

$\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{E[g^2]_t + \epsilon}} g_t$

其中：

$g_t$ 是当前的梯度。
$E[g^2]_t$ 是梯度平方的指数加权平均。
$\eta$ 是学习率。
$\epsilon$ 是为了避免除以零的一个小值（通常取 $10^{-8}$ ）。
$\beta$ 是用于控制移动平均的超参数（通常取值为 0.9）。

1.2 工作机制

RMSProp 通过记录梯度的平方并对其进行加权平均来调节学习率。这样，学习率对频繁更新的参数变小，而对变化不大的参数保持相对较大，从而提高收敛速度，减少训练过程中的振荡。

二、正则化的引入

为了防止模型过拟合，常用的正则化技术包括 L2 正则化（也称为权重衰减）和 RMSProp 的动态调整能力相结合。

2.1 L2 正则化

L2 正则化通过在损失函数中加入一个与参数值相关的惩罚项，限制模型参数的过大增长，避免过拟合。L2 正则化的损失函数形式为：

$L(\theta) = L_{\text{data}}(\theta) + \lambda \|\theta\|_2^2$

其中：

$L_{\text{data}}(\theta)$ 是原始损失函数（如交叉熵损失或均方误差）。
$\lambda$ 是正则化系数，控制正则化强度。
$\|\theta\|_2^2$ 是参数的 L2 范数（权重平方和）。

2.2 RMSProp 与 L2 正则化的结合

在 RMSProp 优化器的基础上，L2 正则化会通过梯度更新过程中引入额外的权重惩罚项，公式如下：

$\theta_t = \theta_{t-1} - \eta \left( \frac{g_t}{\sqrt{E[g^2]_t + \epsilon}} + \lambda \theta_{t-1} \right)$

这里，L2 正则化的影响体现在梯度更新项中，加入了 $\lambda \theta_{t-1}$ ，即对参数 $\theta$ 进行衰减，迫使参数值趋向较小值，以减少模型的复杂度和过拟合风险。

ROPE

在 Transformer 中，由于模型没有卷积和循环结构，因此需要引入位置编码来捕捉序列中词与词之间的相对位置信息。常见的 绝对位置编码（absolute positional encoding） 方式如下：

$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{\frac{2i}{d}}}\right)$

$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{\frac{2i}{d}}}\right)$

其中 pos 是序列中单词的位置，i 是嵌入维度索引，d 是嵌入的维度大小。

局限性：

不能捕捉相对位置信息：绝对位置编码只对词的位置进行编码，不能直接体现词与词之间的相对位置。
对于长序列效果较差：随着序列长度的增加，绝对位置编码容易失效。

二、RoPE 的基本思想

RoPE 的核心思想是通过旋转向量来引入位置信息。具体来说，RoPE 将每个词的嵌入向量根据其在序列中的位置进行旋转变换，从而隐式地引入位置信息，并能够自然地捕捉词与词之间的相对位置关系。

RoPE 的旋转变换方式如下：

通过定义一个二维旋转矩阵，对嵌入向量的每一对维度进行旋转。
旋转的角度与词在序列中的位置有关，因此词的位置信息被编码在旋转变换中。

二维旋转矩阵

RoPE 在嵌入向量的维度上，每两个维度作为一对，通过旋转矩阵对其进行旋转，公式如下：

对于第 i 维度，位置为 pos，角度 $\theta = \frac{pos}{10000^{\frac{2i}{d}}}$ ：

$z' = z \cdot \text{RoPE}(pos)$

其中 $\text{RoPE}(pos)$ 代表旋转变换。该变换通过旋转矩阵的方式作用于每对嵌入维度 (2i,2i+1)：

$\begin{bmatrix} z_{2i}' \\ z_{2i+1}' \end{bmatrix} = \begin{bmatrix} \cos(\theta) & -\sin(\theta) \\ \sin(\theta) & \cos(\theta) \end{bmatrix} \begin{bmatrix} z_{2i} \\ z_{2i+1} \end{bmatrix}$

通过这个旋转操作，嵌入向量的每对维度根据位置 pos 被旋转了不同的角度，从而将位置信息编码到向量中。

2.1 相对位置关系的捕捉

RoPE 的最大优势是能够捕捉词与词之间的相对位置信息，而不需要显式地对相对位置进行编码。对于任意两个词 wiw_iwi 和 wjw_jwj，它们经过 RoPE 编码后的相对位置可以通过旋转向量的差值自然得到。因此，RoPE 能够在长序列中更好地建模词语之间的依赖关系。

三、RoPE 的数学原理

RoPE 将绝对位置信息通过旋转的方式变换到词嵌入中，这种变换具有以下两个重要性质：

相对不变性：RoPE 的旋转编码使得向量之间的相对角度仅与词语的相对位置有关，而不是绝对位置。这意味着，无论句子中的词语在什么位置，它们的相对关系都能被捕捉到。
周期性和旋转不变性：由于 RoPE 的编码基于正弦和余弦函数，编码本身具有周期性。这使得 RoPE 能够在长序列任务中更好地处理序列中词语的位移问题。

四、RoPE 的代码实现

下面是 RoPE 的简单实现，基于 PyTorch：

import torch
import math

def apply_rope(x, seq_len, dim):
    """
    RoPE实现, 将嵌入向量进行旋转变换
    x: 输入的词嵌入向量，形状为 (batch_size, seq_len, dim)
    seq_len: 序列长度
    dim: 嵌入向量维度
    """
    position_ids = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
    theta = torch.arange(0, dim, 2, dtype=torch.float) / dim
    theta = 1.0 / (10000 ** theta)  # 位置编码的角度
    angle = position_ids * theta  # 每个位置对应的角度
    
    # 构造旋转矩阵，将嵌入向量的每对维度旋转
    sin_angle = torch.sin(angle)
    cos_angle = torch.cos(angle)
    
    # 对每对维度进行旋转变换
    x1 = x[..., 0::2] * cos_angle - x[..., 1::2] * sin_angle
    x2 = x[..., 0::2] * sin_angle + x[..., 1::2] * cos_angle
    return torch.cat([x1, x2], dim=-1)

# 示例使用
batch_size = 2
seq_len = 4
dim = 6
x = torch.randn(batch_size, seq_len, dim)

rope_encoded = apply_rope(x, seq_len, dim)
print(rope_encoded)

解释：

apply_rope 函数对输入的词嵌入向量 xxx 进行 RoPE 变换。该向量形状为 (batch_size, seq_len, dim)，表示一个 batch 中多条序列的嵌入。
position_ids 是每个词的位置，theta 是用于旋转的角度。
通过 sin 和 cos 分别对每对维度进行旋转，最后得到经过 RoPE 编码的词嵌入。

Transformer

Transformer是一种基于自注意力机制（Self-Attention Mechanism）的神经网络架构，最初由Vaswani等人在2017年提出，用于解决序列到序列（Sequence-to-Sequence）的问题。Transformer的主要特点是放弃了传统的循环神经网络（RNN）和卷积神经网络（CNN），而是完全依赖于自注意力机制来捕捉输入序列中的长距离依赖关系。

Transformer的核心组件是自注意力层（Self-Attention Layer），它允许模型在不同位置的输入之间建立动态的关联。自注意力层的输出是一个加权和，其中权重是根据输入之间的相似性计算得到的。这使得模型能够关注到与当前位置相关的其他位置的信息。

除了自注意力层，Transformer还包括前馈神经网络（Feed-Forward Neural Network）和残差连接（Residual Connection）等组件。这些组件共同构成了Transformer的编码器（Encoder）和解码器（Decoder）部分，分别用于处理输入序列和生成输出序列。

Llama3 架构

引入库文件

import math
import struct
import inspect
from dataclasses import dataclass
from typing import Any, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

模型参数定义

@dataclass
class ModelArgs:
    dim: int = 4096
    n_layers: int = 6
    n_heads: int = 6
    n_group: Optional[int] = 3
    vocab_size: int = 4096
    hidden_dim: Optional[int] = None
    multiple_of: int = 256  # MLP层隐层维度的指定计算参数(见FFN层)
    norm_eps: float = 1e-5
    max_seq_len: int = 2048
    dropout: float = 0.0

dim: int = 4096: 这是一个整数类型的属性，表示模型的维度，默认值为4096。在深度学习模型中，维度通常指的是输入、输出或中间层的特征数量。

n_layers: int = 6: 表示模型中堆叠的层数，默认为6层。

n_heads: int = 6: 表示多头注意力机制中头的数量，默认为6。在Transformer模型中，多头注意力可以并行处理信息，提高模型的表达能力。

n_group: Optional[int] = 3: 这是一个可选的整数类型属性，表示分组的数量，默认为3。在某些模型中，可能会将数据分组处理以提高效率或性能。

vocab_size: int = 4096: 表示词汇表的大小，默认为4096。在自然语言处理任务中，词汇表包含了模型能够识别的所有词汇。

hidden_dim: Optional[int] = None: 这是一个可选的整数类型属性，表示隐藏层的维度。在某些模型中，这个值可能会被用来指定中间层的大小，但在这里它默认为None，意味着可能在其他地方定义或根据其他参数计算。

multiple_of: int = 256: 这是一个整数类型的属性，用于指定MLP（多层感知机）层隐层维度的计算参数。在设计模型时，有时需要确保某些维度是某个数的倍数，以便于硬件优化或模型设计。

norm_eps: float = 1e-5: 这是一个浮点数类型的属性，表示归一化时的epsilon值，默认为0.00001。在归一化操作中，epsilon用于防止除以零的情况发生。

max_seq_len: int = 2048: 表示模型能够处理的最大序列长度，默认为2048。在处理文本数据时，这个参数限制了模型能够处理的最长文本长度。

dropout: float = 0.0: 表示在训练过程中随机丢弃（dropout）的比率，默认为0.0，即不进行dropout。Dropout是一种正则化技术，用于防止模型过拟合。

RMS正则化

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps
        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

class RMSNorm(torch.nn.Module): 这行代码声明了一个继承自torch.nn.Module的类RMSNorm。torch.nn.Module是所有神经网络模块的基类，提供了一些基本的功能，如参数管理和前向传播。

def __init__(self, dim: int, eps: float): 这是类的构造函数，用于初始化RMSNorm层。它接受两个参数：

dim: int: 表示输入特征的维度。
eps: float: 一个很小的正数，用于数值稳定性，防止除以零。

super().__init__(): 这行代码调用父类torch.nn.Module的构造函数，完成基本的初始化。

self.eps = eps: 将传入的eps值保存为类的属性。

self.weight = nn.Parameter(torch.ones(dim)): 创建一个参数weight，它是一个形状为(dim,)的一维张量，初始值为1。nn.Parameter是PyTorch中用于定义可学习的参数的类。

def _norm(self, x): 这是一个辅助函数，用于计算RMSNorm的归一化操作。它接受一个输入张量x，并返回归一化后的张量。

x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps): 这是RMSNorm的核心计算步骤。
x.pow(2): 计算输入x的元素平方。
.mean(-1, keepdim=True): 计算最后一个维度（通常是特征维度）的均值，并保持维度不变。
torch.rsqrt(...): 计算倒数平方根（即平方根的倒数）。
+ self.eps: 加上一个小的正数eps，以提高数值稳定性。

def forward(self, x): 这是前向传播函数，用于定义模块如何处理输入数据并产生输出。

output = self._norm(x.float()).type_as(x): 首先将输入x转换为浮点数类型，然后调用_norm函数进行归一化。type_as(x)确保输出张量与输入张量具有相同的数据类型。
return output * self.weight: 最后，将归一化后的输出乘以权重weight，得到最终的输出。

ROPE相关

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device) 
    freqs = torch.outer(t, freqs).float() 
    freqs_cos = torch.cos(freqs) 
    freqs_sin = torch.sin(freqs) 
    return freqs_cos, freqs_sin

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(shape)

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cos: torch.Tensor,
    freqs_sin: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:

    # 重塑 xq 和 xk，使其与复数表示相匹配
    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)

    # 重塑形为了广播
    freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
    freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)

    # 应用旋转嵌入
    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos

    # 讲最后两维度拉平。
    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

1. precompute_freqs_cis 函数

这个函数用于预先计算频率和余弦正弦值，这些值将用于旋转位置编码。

参数:

dim: 整数，表示维度。
end: 整数，表示序列的长度。
theta: 浮点数，用于计算频率的参数，默认值为10000.0。

过程:

使用torch.arange生成一个从0开始到dim的序列，步长为2，然后除以dim并取其倒数，再对结果进行theta次方的倒数，得到频率。
使用torch.arange(end, device=freqs.device)生成一个从0到end-1的序列，用于与频率进行外积运算，得到每个位置的频率。
计算外积的结果的余弦和正弦值。

返回值:

返回两个张量：freqs_cos和freqs_sin，分别包含所有位置的余弦和正弦值。

2. reshape_for_broadcast 函数

这个函数用于调整频率张量的形状，使其可以与输入张量进行广播操作。

参数:

freqs_cis: 频率张量。
x: 输入张量。

过程:

检查freqs_cis的形状是否与x的第二和最后一维匹配。
根据x的维度生成一个新的形状，其中第二和最后一维保持不变，其他维度设置为1，以便进行广播。

返回值:

返回重新形状后的频率张量。

3. apply_rotary_emb 函数

这个函数用于应用旋转位置编码到输入张量。

参数:

xq, xk: 输入张量，通常代表查询（query）和键（key）。
freqs_cos, freqs_sin: 余弦和正弦频率张量。

过程:

将xq和xk重塑并分解为实部和虚部。
使用reshape_for_broadcast调整频率张量的形状。
应用旋转公式，其中实部和虚部分别乘以余弦和正弦值，然后进行适当的加减操作。
将结果重新堆叠并拉平。

返回值:

返回旋转编码后的xq和xk张量。

# 定义输入x， n_rep是需要重复的次数，在这里一般是组数
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:

    bs, slen, n_kv_heads, head_dim = hidden_states.shape
    # dont need repeat here means multi head attention
    if n_rep == 1:
        return hidden_states
    # first we expand x to (bs, seq_len, head, group, head_dim)
    hidden_states = hidden_states[:, :, :, None, :].expand(bs, slen, n_kv_heads, n_rep, head_dim)
    # reshape make head -> head * group
    return hidden_states.reshape(bs, slen, n_kv_heads * n_rep, head_dim)

这段代码是一个Python函数，它定义了一个名为 repeat_kv 的函数，这个函数接受两个参数：hidden_states 和 n_rep。hidden_states 是一个PyTorch张量（tensor），它代表神经网络中的隐藏状态；n_rep 是一个整数，表示需要重复的次数，通常用于多头注意力机制中的组数。

函数的主要步骤如下：

首先，函数获取 hidden_states 张量的形状，这包括批次大小（bs）、序列长度（slen）、键值头数（n_kv_heads）和每个头的维度（head_dim）。

如果 n_rep 等于 1，这意味着不需要重复，函数直接返回原始的 hidden_states 张量。

如果需要重复（即 n_rep 大于 1），函数将 hidden_states 张量扩展到一个新的形状 (bs, slen, n_kv_heads, n_rep, head_dim)。这是通过在 hidden_states 张量中添加一个新的维度并使用 expand 方法来实现的。expand 方法不会实际复制数据，而是创建一个新的视图，其中重复了原始数据。

然后，函数将扩展后的张量重新塑形（reshape），将头数（n_kv_heads）和重复次数（n_rep）相乘，得到一个新的头数。这样，每个头的维度（head_dim）保持不变。

最后，函数返回重新塑形后的张量。

Attention

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        
        self.group = args.n_group
        self.heads = args.n_heads
        self.kv_heads = args.n_heads // args.n_group
        assert args.n_heads % self.kv_heads == 0
        self.head_dim = args.dim // args.n_heads
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)
        self.dropout = args.dropout
        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
        mask = torch.triu(mask, diagonal=1)
        self.register_buffer("mask", mask)

    def forward(
        self,
        x: torch.Tensor,
        freqs_cos: torch.Tensor,
        freqs_sin: torch.Tensor,
    ):
        bsz, seqlen, _ = x.shape

        # QKV
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        xq = xq.view(bsz, seqlen, self.heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.kv_heads, self.head_dim)

        # RoPE relative positional embeddings
        xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)

        # grouped multiquery attention: expand out keys and values
        xk = repeat_kv(xk, self.group)  # (bs, seqlen, n_local_heads, head_dim)
        xv = repeat_kv(xv, self.group)  # (bs, seqlen, n_local_heads, head_dim)

        # make heads into a batch dimension
        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)



        # 先不使用flash attn，从零走一遍流程！
        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
        assert hasattr(self, 'mask')
        scores = scores + self.mask[:, :, :seqlen, :seqlen]   # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        scores = self.attn_dropout(scores)
        output = torch.matmul(scores, xv)  # (bs, n_local_heads, seqlen, head_dim)

        # restore time as batch dimension and concat heads
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

        # 最终送入output层并正则，得到最终结果。
        output = self.wo(output)
        output = self.resid_dropout(output)
        return output

构造函数 __init__:

接受一个 ModelArgs 类型的参数 args，这个参数包含了模型的各种配置。
初始化分组数 self.group，头数 self.heads，键值头数 self.kv_heads，以及每个头的维度 self.head_dim。
定义线性层 self.wq, self.wk, self.wv, self.wo 分别用于计算查询（Q）、键（K）、值（V）和输出（O）。
初始化 Dropout 层 self.attn_dropout 和 self.resid_dropout。
创建一个上三角掩码 mask 并注册为缓冲区，这个掩码用于在自注意力计算中防止位置信息的泄露。

前向传播函数 forward:

接受输入张量 x 和两个频率张量 freqs_cos, freqs_sin，这些张量用于计算相对位置编码。
计算查询（Q）、键（K）、值（V）并通过线性层进行变换。
将 Q、K、V 重塑为适合多头注意力的形状。
应用相对位置编码（RoPE）到 Q 和 K。
使用 repeat_kv 函数扩展 K 和 V，以适应分组的多头注意力。
将 Q、K、V 的头维度转换为批次维度。
计算 Q 和 K 的点积，然后应用掩码和 softmax 函数来获得注意力分数。
将注意力分数应用到 V 上，得到输出。
将输出的头维度合并，并应用输出线性层和残差连接的 Dropout。
返回最终的输出张量。

这个注意力机制的特点是：

使用分组多头注意力，其中每个组内的头共享相同的键和值。
引入了相对位置编码（RoPE），这是一种编码序列中相对位置信息的方法，可以增强模型对序列顺序的理解。
使用 Dropout 来防止过拟合，并在训练过程中增加正则化。

FFN网络

class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = 4 * dim
            hidden_dim = int(2 * hidden_dim / 3)
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // 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)
        self.dropout = nn.Dropout(dropout)

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

构造函数 __init__:

接受参数 dim（输入和输出维度）、hidden_dim（中间层维度）、multiple_of（维度对齐因子）、dropout（Dropout 比率）。
如果 hidden_dim 没有提供（即 None），则使用默认的计算方式来确定它的值。这个默认值是输入维度的四倍，然后通过一个转换公式来确保它是 multiple_of 的整数倍。
定义三个线性层 self.w1、self.w2 和 self.w3。self.w1 将输入维度映射到 hidden_dim，self.w2 将 hidden_dim 映射回 dim，而 self.w3 再次将 dim 映射到 hidden_dim。
初始化 Dropout 层 self.dropout。

前向传播函数 forward:

接受输入张量 x。
首先，x 通过 self.w1 线性层。
然后，使用 SiLU（Sigmoid线性单元）激活函数（也称为 Swish）处理 self.w1(x) 的输出。
接着，将 self.w3(x) 的结果与 SiLU 激活后的结果相乘。
最后，将乘积通过 self.w2 线性层，并通过 self.dropout 应用 Dropout，然后返回最终的输出。

这个前馈网络的特点是：

使用了 SiLU 激活函数，这是一种平滑的非单调激活函数，有助于模型的训练和泛化。
通过 self.w3 引入了一个额外的线性变换，这在某些情况下可以增加模型的表达能力。
使用 Dropout 来防止过拟合，并在训练过程中增加正则化。

decoder-layer

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=args.hidden_dim,
            multiple_of=args.multiple_of,
            dropout=args.dropout,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(self, x, freqs_cos, freqs_sin):
        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
        out = h + self.feed_forward.forward(self.ffn_norm(h))

构造函数 __init__:

接受参数 layer_id（表示当前层的编号）和 args（包含模型配置的 ModelArgs 类型）。
初始化头数 self.n_heads，维度 self.dim，以及每个头的维度 self.head_dim。
创建一个 Attention 实例 self.attention，用于实现自注意力机制。
创建一个 FeedForward 实例 self.feed_forward，用于实现前馈网络。
初始化两个 RMSNorm 实例 self.attention_norm 和 self.ffn_norm，用于层归一化。

前向传播函数 forward:

接受输入张量 x 和两个频率张量 freqs_cos, freqs_sin，这些张量用于相对位置编码。
首先，将输入 x 通过层归一化 self.attention_norm，然后传递给自注意力层 self.attention。
将自注意力层的输出与输入 x 相加，实现残差连接。
将残差连接的结果传递给另一个层归一化 self.ffn_norm，然后传递给前馈网络 self.feed_forward。
将前馈网络的输出与残差连接的结果相加，再次实现残差连接。
返回最终的输出张量。

这个变换器块的特点是：

使用了残差连接和层归一化，这有助于改善训练过程中的梯度流动，防止梯度消失或爆炸问题。
自注意力机制允许模型在序列的不同位置之间动态地分配不同的注意力权重。
前馈网络提供了额外的非线性变换能力。

class Transformer(nn.Module):
    last_loss: Optional[torch.Tensor]

    def __init__(self, params: ModelArgs):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers

        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)  # 其weight形状为(vocab,dim)
        self.dropout = nn.Dropout(params.dropout)
        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)  # 维数也为(vocab,dim)--x·W^T

        # 将模型的嵌入层（embedding layer）和输出层（unembedding layer）的权重共享，即 "权重共享" 或 "weight tying"
        self.tok_embeddings.weight = self.output.weight # 来源论文: https://paperswithcode.com/method/weight-tying

        # some useful precompute for the RoPE relative positional embeddings
        freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

        # init all weights
        self.apply(self._init_weights)
        # apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers))

        # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor.
        self.last_loss = None

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None) -> torch.Tensor:
        _bsz, seqlen = tokens.shape
        h = self.tok_embeddings(tokens)
        h = self.dropout(h)
        freqs_cos = self.freqs_cos[:seqlen]
        freqs_sin = self.freqs_sin[:seqlen]

        for layer in self.layers:
            # print('loging')
            h = layer(h, freqs_cos, freqs_sin)
        h = self.norm(h)

        if targets is not None:
            # 有targets则计算loss
            logits = self.output(h)
            self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # 在推理阶段，只抽取最后一行--预测下一个token即可
            logits = self.output(h[:, [-1], :]).reshape(_bsz,-1) # note: using list [-1] to preserve the time dim
            self.last_loss = None

        return logits

构造函数 __init__:

接受参数 params，它是一个包含模型配置的 ModelArgs 类型。
初始化词汇表大小 self.vocab_size 和层数 self.n_layers。
创建一个嵌入层 self.tok_embeddings，用于将词汇表中的单词映射到维度空间。
创建一个 Dropout 层 self.dropout。
创建一个 TransformerBlock 层的列表 self.layers，每个层对应一个 TransformerBlock 实例。
创建一个层归一化层 self.norm。
创建一个输出层 self.output，用于将变换器的输出映射回词汇表大小的维度空间。
实现权重共享，将嵌入层和输出层的权重设置为相同。
为相对位置编码预计算频率值，并将它们注册为缓冲区。
初始化所有权重，包括特殊的缩放初始化，用于残差投影。
初始化一个属性 self.last_loss 用于存储最后一次前向传播的损失。

前向传播函数 forward:

接受输入张量 tokens 和可选的目标张量 targets。
计算嵌入层的输出 h 并应用 Dropout。
根据序列长度截取相对位置编码的频率值。
通过每个 TransformerBlock 层处理 h。
应用层归一化。
如果提供了 targets，则计算交叉熵损失并将 self.last_loss 设置为该损失值。
如果没有提供 targets，则提取 h 的最后一行作为预测下一个 token 的 logits。
返回 logits。

这个变换器模型的特点是：

使用权重共享来减少模型参数的数量，提高模型的性能。
通过相对位置编码（RoPE）来引入序列中的位置信息。
使用残差连接和层归一化来改善训练过程中的梯度流动。
可以用于语言模型或其他需要处理序列数据的任务。