Meta最新模子LLaMA细节与代码详解

0. 简介

本日先容的内容是Facebook Meta AI最新提出的语言模子LLaMA，该模子声称以更小的体积，在多数任务上逾越了GPT-3的性能。
模子相干项目已经开源：
https://github.com/facebookresearch/llama
论文所在：https://scontent-tpe1-1.xx.fbcdn.net/v/t39.8562-6/333078981_693988129081760_4712707815225756708_n.pdf?_nc_cat=108&ccb=1-7&_nc_sid=ad8a9d&_nc_ohc=ov6yTHfLfNQAX-guxqd&_nc_ht=scontent-tpe1-1.xx&oh=00_AfDMyTEYewg-cHT9_4_sUaW5h0gqrqwjcNMylD9HtVFCWA&oe=6401C9E2
由于模子较大，现在的装备临时没有办法支持进一步的实验，但是其模子代码已经开源，以是可以先通过代码相识一下模子结构上的一些细节，本日就针对github上放出的代码，相识一下模子的细节。
别的，该模子实在就是transformer做了一点细节上的改进，真正更有代价的工作应该在数据和训练方面。通过阅读代码，可以对transformer的底子构造举行复习，而且相识大模子怎样在多卡上分布推理。
由于该项目源码险些没有表明，这就肯定会给许多同砚阅读时带来困扰，以是本文顺带着就把代码部门详细的先容一下。
1. 项目情况依赖

此项目给出的情况依赖只有4个：

  

• torch  
  
• fairscale  
  
• fire  
  
• sentencepiece
  

此中torch不比多讲，fairscale是用来做GPU分布的，一样寻常是当利用DDP仍旧碰到超显存的题目时利用fairscale。现在fairscale我还没有试过，在下文的源码先容中，我会用torch中对应的底子网络替换fairscale中的结构层举行先容。fire是一个下令行工具，用大概不消他都可以，sentencepiece是用于tokenizer的工具包，会在tokenizer部门简单先容。
2. 模子细节

由于该模子就是用的transformer的decoder，以是在结构上它与GPT好坏常雷同的，只是有一些细节必要注意一下。
2.1 RMS Pre-Norm

关于Pre-Norm和Post-Norm是神经网络中老生常谈的话题，现在比力广泛的被各人继承的结论是，雷同的深度条件下，Post-Norm的结果要优于Pre-Norm，由于Pre-Norm现实上相称于通过了一个更宽的网络而非更深的网络，以是在划一深度下，Pre-Norm的现实结果相称于一个更浅却更宽的网络，详细的推理过程参考：https://spaces.ac.cn/archives/9009。
然而在LLaMA中却接纳了Pre-Norm，大概是由于模子够深（7B，13B，30B，65B的模子，transformer layer数目分别为32，40，60，80），而Pre-Norm的恒中分支更加显着，有利于梯度的传播（这部门临时没有想到很公道的表明，如果有更好的明白，欢迎在批评区增补）。
RMS Norm（Root Mean Square Layer Normalization），是一样寻常LayerNorm的一种变体，可以在梯度降落时令丧失更加平滑。
与layerNorm相比，RMS Norm的重要区别在于去掉了减去均值的部门（re-centering），只保存方差部门（re-scaling），从归一化的表达式上可以直观地看出：

  

• 一样寻常的LN：
  

  a ‾ i = a i − μ σ g i \overline{a}_i = \frac {a_i- \mu} \sigma g_i ai​=σai​−μ​gi​
此中，
  μ = 1 n ∑ i = 1 n a i \mu = \frac 1 n \sum_{i=1}^na_i μ=n1​i=1∑n​ai​
  σ = 1 n ∑ i = 1 n ( a i − μ ) 2 \sigma= \sqrt {\frac 1 n \sum_{i=1}^n{      {(a_i-\mu)}^2}} σ=n1​i=1∑n​(ai​−μ)2             ​

  

• RMS Norm：
    a ‾ i = a i R M S ( a ) g i \overline{a}_i = \frac {a_i} {RMS(a)} g_i ai​=RMS(a)ai​​gi​
  此中，
    R M S ( a ) = 1 n ∑ i = 1 n a i 2 {RMS(a)}=\sqrt {\frac 1 n \sum_{i=1}^n{       {a_i}^2}} RMS(a)=n1​i=1∑n​ai​2              ​
  

可以看到，二者的区别就在于有没有减去均值。至于RMS Norm为什么有效，必要求梯度举行分析，感爱好的同砚可以阅读RMS Norm的论文。
2.2 SwiGLU激活函数

LLaMA接纳SwiGLU更换了原有的ReLU。
接纳SwiGLU的FNN，在论文中以如下公式举行表述：
  F F N s w i G L U ( x , W , V , W 2 ) = ( S w i s h 1 ( x W ) ⊗ x V ) W 2 FFN_{swiGLU}(x, W, V, W_2) = (Swish_1(xW)\otimes xV)W_2 FFNswiGLU​(x,W,V,W2​)=(Swish1​(xW)⊗xV)W2​
此中， S w i s h β ( x ) = x σ ( β x ) Swish_\beta(x) = x\sigma(\beta x) Swishβ​(x)=xσ(βx), (Ramachandran et al., 2017.)
2.3 RoPE旋转位置编码

RoPE（Rotary Position Embedding）旋转位置编码，是苏剑林老师提出的一种旋转位置编码方法，其头脑是接纳绝对位置编码的情势，实现相对位置编码。这一部门比力关键，如果不明白的话，后边的代码估计就看不懂了。读懂RoPE涉及一点复变函数的底子知识，不外如果你没有学过的话也没有关系。
位置编码对大模子而言尤为告急，由于既然是要训练大模子，那么长文本的表征和模子对于长文本的建模本领就显得非常告急。（但是对于绝对位置编码，我有一个直观地感受，以为其本质上不实用于长文本的场景，由于它会直接导致模子的Embedding层被无穷放大，而且由于数据分布在seq_len方向上通常是长尾的，这又会肯定导致绝对位置编码的矩阵在尾部会越来越希罕，一方面造成资源浪费，另一方面这种表现方法直观上就很倒霉于模子的学习，由于它与我们现实场景是有很大的抵牾的。而RoPE固然具有相对位置编码的性子，但是从代码部门可以看出，在构造的时间，其也是受到了最大长度的限定的。关于这一点，我无法严谨得阐明，只是一点个人的想法。）。
而RoPE的奥妙之处在于，它既保存了绝对位置编码中的绝对位置信息，又保存了在内积运算下，对位置信息的相对性。
RoPE重要借助了复数的头脑。为了引入复数，起首假设了在到场位置信息之前，原有的编码向量是二维行向量 q m q_m qm​和 k n k_n kn​，此中 m m m和 n n n是绝对位置，现在必要构造一个变更，将 m m m和 n n n引入到 q m q_m qm​和 k n k_n kn​中，即探求变更：
  q m ~ = f ( q , m ) , k n ~ = f ( k , n ) \tilde {q_m} = f(q, m), \tilde{k_n} = f(k, n) qm​~​=f(q,m),kn​~​=f(k,n)
思量到Attention的核心盘算是内积：
  A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V Attention(Q, K,V) = softmax(\frac {QK^T} {\sqrt{d_k}})V Attention(Q,K,V)=softmax(dk​                    ​QKT​)V
以是，寻求的这个 f ( ∗ ) f(*) f(∗)变更，应该具有特性： ⟨ f ( q , m ) , f ( k , n ) ⟩ = g ( q , k , m − n ) \langle f(q, m), f(k, n) \rangle = g(q, k, m-n) ⟨f(q,m),f(k,n)⟩=g(q,k,m−n)
这里直接说结论，寻求的变更就是 q m e i m θ q_me^{im\theta} qm​eimθ，也就是给 q m q_m qm​乘以 e i m θ e^{im\theta} eimθ，相应地， k n k_n kn​乘以 e i n θ e^{in\theta} einθ。
详细的求解过程，请参考苏剑林老师的博客。
做了如许一个变更之后，根据复数的特性，有：
  ⟨ q m , k n ⟩ = R e [ q m k n ∗ ] \langle q_m, k_n \rangle = Re[q_mk^*_n] ⟨