llama3 结构详解

1. Llama3 整体结构

  llama3 的整体结构还是延续transformer decoder 架构，其整体架构如下图左侧蓝色虚线框中所示。模型结构并不复杂，其主要组件为32个Transformer Block(32 为meta llama3 中的默认值)(见下图红色虚线框中所示)。

${注}_1$： 下一节中会参照上图中 红色圆形序号 讲解各模块。
${注}_2$： llama3的RoPE算法被拆成了3个方法来实现，上图中的模块2只包含了一个方法，另两个方法是在Attention模块(模块5)中进行的调用。

2. 模块详解

2.1 模块1: Embeddings

  llama3 的embedding 使用的是VocabParallelEmbedding这个类进行的向量转换，这个类是meta的fairscale包中的一个类，可以理解为对torch.nn.embedding做了并行化。

2.2 模块2: RoPE

  前文中已经提及llama3的RoPE算法被拆成了3个方法来实现，模块2只包含了一个方法，另两个方法是在Attention模块(模块5)中进行的调用。本小节具体按照RoPE的原始论文来讲解，主要阐述RoPE的算法原理。

2.2.1 从一个2维的例子说起 RoPE

  我们知道，寻找位置编码的基本思路是 输入位置编码经过特征提取的核心算法后的值，应能反应出两个位置之间的先后顺序(这点不是必要的)和相对位置信息。（《Transformer(二)–论文理解：transformer 结构详解》 2.1节 中有简单说明），RoPE的原始论文中给出了一个数学表达，如下式：
$$<f_q(x_m,m),f_k(x_n,n)>=g(x_m,x_n,m-n)\text{\hspace{1em}(2.1)}$$

  $f_q(x_m,m)$和$f_k(x_n,n)$分别表示$q_m$和$k_n$。式子的左侧为点积形式，之所以为点积是因为tansformer中使用的attention score计算方法通常为点积。右侧$g(x_m,x_n,m-n)$表示计算结果是与$x_m,x_n,m-n$相关的。在这里$m-n$的绝对值能反应出位置的距离，大小反应出前后顺序。 我们的目的就是找到一个这样的变换函数 f { q , k } f_{\{q,k\}} f{q,k}​能表达$f_q(x_m,m)$与$f_k(x_n,n)$，使$f_q$与$f_k$做点积操作后能保留$m-n$的信息。当然我们找到了，见公式2.2

  RoPE的论文中是先从2D情况下举例说明我们找到的$f(x)$的，如下，当$d=2$时：

$$\begin{gathered} f_q(x_m,m)=(W_q{x}_m)e^{im\theta } \\ {f}_k(x_n,n)=(W_k{x}_n)e^{in\theta } \\ g(x_m,x_n,m-n)=Re[(W_q{x}_m)(W_k{x}_n{)}^{\ast }{e}^{i(m-n)\theta }]\text{\hspace{1em}(2.2)} \end{gathered}$$

  其中$Re[\cdot ]$是复数的实部，$(W_k{x}_n{)}^{\ast }$表示$(W_k{x}_n)$的共轭复数。
$\theta \in \mathbb{R}$ 是一个预设的非零常数。我们可以进一步将 f { q , k } f_{\{q,k\}} f{q,k}​写成乘法矩阵：
f { q , k } ( x m , m ) = ( c o s   m θ − s i n   m θ s i n   m θ c o s   m θ ) ( W { q , k } ( 11 ) W { q , k } ( 12 ) W { q , k } ( 21 ) W { q , k } ( 22 ) ) ( x m ( 1 ) x m ( 2 ) ) (2.3) f_{\{q,k\}}(x_m,m)= \left( \begin{matrix} cos\ m\theta & -sin\ m\theta \\ sin\ m\theta & cos\ m\theta \\ \end{matrix} \right) \left( \begin{matrix} W^{(11)}_{\{q,k\}} & W^{(12)}_{\{q,k\}} \\ W^{(21)}_{\{q,k\}} & W^{(22)}_{\{q,k\}} \\ \end{matrix} \right) \left( \begin{matrix} x^{(1)}_{m} \\ x^{(2)}_{m} \end{matrix} \right) \tag{2.3} f{q,k}​(xm​,m)=(cos mθsin mθ​−sin mθcos mθ​)(W{q,k}(11)​W{q,k}(21)​​W{q,k}(12)​W{q,k}(22)​​)(xm(1)​xm(2)​​)(2.3)

其中，$(x_m^{(1)},x_m^{(2)})$是$x_m$在二维坐标系中的表示。同样的，$g$也可以看作一个矩阵，因此可以在2维情况下求解公式(2.1)。

2.2.2 RoPE的一般形式

  为了将我们在2D中的结果推广到任意的$x_i\in {\mathbb{R}}^d$，我们将d维空间划分为d/2个子空间，并根据内积的线性性质将它们组合起来，将 f { q , k } ( x m , n ) f_{\{q,k\}}(x_m,n) f{q,k}​(xm​,n)转化为：
f { q , k } ( x m , m ) = R Θ , m d W { q , k } x m (2.4) f_{\{q,k\}}(x_m,m)=\pmb{R}^{d}_{\Theta, m}\pmb{W}_{\{q,k\}}x_m \tag{2.4} f{q,k}​(xm​,m)=RΘ,md​W{q,k}​xm​(2.4)

  其中， W { q , m } \pmb{W}_{\{q,m\}} W{q,m}​ 表示与query和key 所对应的转换矩阵 ，$x_m$ 为输入向量，$R_{\Theta ,m}^d$为旋转矩阵，具体如下：
$$R_{\Theta ,m}^d=\left(\begin{array}{ccccccc}cosm{\theta }_1& -sinm{\theta }_1& 0& 0& \cdots & 0& 0\\ sinm{\theta }_1& cosm{\theta }_1& 0& 0& \cdots & 0& 0\\ 0& 0& cosm{\theta }_2& -sinm{\theta }_2& \cdots & 0& 0\\ 0& 0& sinm{\theta }_2& cosm{\theta }_2& \cdots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \cdots & cosm{\theta }_{d/2}& -sinm{\theta }_{d/2}\\ 0& 0& 0& 0& \cdots & sinm{\theta }_{d/2}& cosm{\theta }_{d/2}\end{array}\right)\text{\hspace{1em}(2.5)}$$

Θ = { θ i = 1000 0 − 2 ( i − 1 ) / d , i ∈ [ 1 , 2 , . . . , d / 2 ] } (2.6) \Theta=\{ \theta_i = 10000^{-2(i-1)/d}, i \in [1,2,...,d/2] \} \tag{2.6} Θ={θi​=10000−2(i−1)/d,i∈[1,2,...,d/2]}(2.6)

2.2.3 RoPE的理解

  这里我们把我们求出的 f { q , k } ( x m , m ) = R Θ , m d W { q , k } x m f_{\{q,k\}}(x_m,m)=\pmb{R}^{d}_{\Theta, m}\pmb{W}_{\{q,k\}}x_m f{q,k}​(xm​,m)=RΘ,md​W{q,k}​xm​代入attention score的计算公式
$$a_{m,n}=\frac{\exp (\frac{q_m^T{k}_n}{\sqrt{d}})}{{\sum }_{j=1}^N\exp (\frac{q_m^T{k}_j}{\sqrt{d}})}\text{\hspace{1em}(2.7)}$$

这里我们只需要看$q_m^T{k}_m$即可，公式的其余部分不会改变结果形式。把公式2.4代入2.7

$$q_m^T{k}_n=(R_{\Theta ,m}^d{W}_q{x}_m{)}^T(R_{\Theta ,n}^d{W}_k{x}_n)=x^T{W}_q{R}_{\Theta ,n-m}^d{W}_k{x}_n\text{\hspace{1em}(2.8)}$$

其中，$R_{\Theta ,n-m}^d=(R_{\Theta ,m}^d{)}^T{R}_{\Theta ,n}^d$，注意 $R_{\Theta }^d$是一个正交矩阵，这保证了位置信息在处理过程中的稳定性。此外，由于$R_{\Theta }^d$的稀疏性，式(2.8)的计算效率不高，作者在理论上提供了另一种实现。

2.3 模块3: Transformer Block

  Transformer Block 模块是llama3的核心模块，或者说，llama3为Transformer Block模块堆叠而成。Transformer Block有模块4、5、6、7组成，具体内容见对应模块。

2.4 模块4: RMSNorm

  RSMNorm 是在 layer normalization 基础上优化而来，所以先简单回顾下layer normalization。（详细介绍见《Transformer(二)–论文理解：transformer 结构详解》 2.4节）
  layer normalization 是根据下面的公式对$x$的分布进行调整。
$$x=a\ast \frac{x-\overline{x}}{std+eps}+b\text{\hspace{1em}(2.9)}$$
其中，$\overline{x}$是均值，$std$是标准差，$eps$为一个很小的数，防止分母为零。$a$、$b$为参数，$b$可以为零。
  我们现在来看看RMSNorm做了什么优化呢，其实他对上面的试子 $x=a\ast \frac{x-\overline{x}}{std+eps}+b$进行了简化。RMSNorm的计算公式如下：
$${\overline{x}}_i=\frac{x_i}{RMS(x)}{g}_i,whereRMS(x)=\sqrt{\frac{1}{n}{\Sigma }_{i=1}^n{x}_i^2}\text{\hspace{1em}(2.10)}$$

  上式中$g_i$为权重参数，可以看出，RMSNorm移除了LayerNorm中的均值项（原式中的$\overline{x}$项），$std$的计算中，也没有做减去均值的操作($std=\sqrt{\frac{1}{n}{\Sigma }_{i=1}^n(x_i-\overline{x})}$)。这种简化在计算效率上有一定提高，且原始论文也说了，在效果上没有明显影响。

下面附上meta llama3中RMSNorm的源码，方便大家理解。

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

2.5 模块5: Attention

  llama3的attention模块主要做了4部分工作，分别是RoPE计算、分注意力分组机制实现、点积注意力计算 及 kv缓存策略实现。其中RoPE的计算在模块2中已经讲解，这里不在赘述。下文对GQA，点积注意力计算及KV缓存进行简单的讲解。

2.5.1 分组注意力机制(GQA)

  llama3中的attention模块与《Attention is all you need》中使用的attention技术有些许优化。同样是使用Scaled Dot-Product Attention来计算attention score，但分组优化这块没有延续使用MHA（Multi-head Attention)技术，而是使用了GQA(Grouped-Query Attention)分组技术。具体的Scaled Dot-Product Attention 与MHA我之前在《Transformer(二)–论文理解：transformer 结构详解》一文的2.2节中，已经写的非常详细了，所以这里不再展开，只讲解下GQA。

  我们知道，在MHA中，由于每个head都有独立的键和值，内存和计算成本较高，特别是在处理长序列或大批量数据时。然后就有大牛Noam Shazeer提出了MQA（Multi Query Attention）方法，将原来的h个KV对缩减为1个，所有query只使用一个共享的KV对，这种改造虽然大大减少了显存消耗，但其特征捕捉能力也受到影响。因此又提出了GQA（Grouped-Query Attention ）， 将query 进行分组，每组共享一个KV对。下面是GQA原始论文中给出的对比图。

2.5.2 注意力计算(Scaled Dot-Product Attention)

  llama3 计算attention score时，使用了与《attention is all you need》一文中相同的计算方法，即点积注意力方法（Scaled Dot-Product Attention）,由于Scaled Dot-Product Attention在《Transformer(二)–论文理解：transformer 结构详解》 一文中的2.2.1章节有详细的讲解，这里就不再展开。

2.5.3 KV缓存

   llama3在计算 attention 时采用了kv cache策略。此策略的思想是缓存每个时间步的key和value的值，在推理阶段，由于模型是自回归模式生成文本，所以当我们对过往时间步有缓存结果时，会减少计算量，提高解码效率。

下面是llama3中Attention类的源码，大家可以参考理解

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
		.
		.
		.
    

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        # 以下是Scaled Dot-Product Attention的计算
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output)

2.6 模块6: ADD

   此模块做了个类似残差的操作，但与残差不同的是，不是用输入减去输出，而是用输入加上输出。具体操作就是把模块4的输入与模块5的输出做加法运算。

2.7 模块7: FFN

  由3个Linear组成的FeedForward网络，这里的激活函数使用的siLU。siLU的数学公式如下：
$$silu(x)=x\ast \sigma (x),where\sigma (x)isthelogisticsigmoid.$$

函数的激活曲线如下图：

在里注意下，siLU 还有一个名字叫“swish function”，这个在 pytorch 的官方文档中有说明。

下面给出主要源码。

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        .
        .
        .
  

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

2.8 模块8: Linear

  此模块的目的是把模型中 decoder的输出从$d_{model}$维度映射到词表大小的维度。下面是meta llama中的linear层的初始化。

 self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )