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

Posted 2023-03-29 常鸿宇

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

• 0. 简介
• 1. 项目环境依赖
• 2. 模型细节
• • 2.1 RMS Pre-Norm
  • 2.2 SwiGLU激活函数
  • 2.3 RoPE旋转位置编码
• 3. 代码解读
• • 3.1 tokenizer
  • 3.2 model
  • • 3.2.1 模型细节详解
    • 3.2.2 transformer构建
  • 3.3 generate
• 4. 推理

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：

$${\overline{a}}_i=\frac{a_i-\mu }{\sigma }{g}_i$$
其中，

$$\mu =\frac{1}{n}\sum _{i=1}^n{a}_i$$
$$\sigma =\sqrt{\frac{1}{n}\sum _{i=1}^n{(a_i-\mu )}^2}$$

• RMS Norm：
  $${\overline{a}}_i=\frac{a_i}{RMS(a)}{g}_i$$
  其中，
  $$RMS(a)=\sqrt{\frac{1}{n}\sum _{i=1}^n{a_i}^2}$$

可以看到，二者的区别就在于有没有减去均值。至于RMS Norm为什么有用，需要求梯度进行分析，感兴趣的同学可以阅读RMS Norm的论文。

2.2 SwiGLU激活函数

LLaMA采用SwiGLU替换了原有的ReLU。

采用SwiGLU的FNN，在论文中以如下公式进行表述：
$$FF{N}_{swiGLU}(x,W,V,W_2)=(Swis{h}_1(xW)\otimes xV)W_2$$

其中，$Swis{h}_{\beta }(x)=x\sigma (\beta x)$, (Ramachandran et al., 2017.)

2.3 RoPE旋转位置编码

RoPE（Rotary Position Embedding）旋转位置编码，是苏剑林老师提出的一种旋转位置编码方法，其思想是采用绝对位置编码的形式，实现相对位置编码。这一部分比较关键，如果不理解的话，后边的代码估计就看不懂了。读懂RoPE涉及一点复变函数的基础知识，不过如果你没有学过的话也没有关系。

位置编码对大模型而言尤为重要，因为既然是要训练大模型，那么长文本的表征和模型对于长文本的建模能力就显得非常重要。（但是对于绝对位置编码，我有一个直观地感受，认为其本质上不适用于长文本的场景，因为它会直接导致模型的Embedding层被无限放大，并且由于数据分布在seq_len方向上通常是长尾的，这又会必然导致绝对位置编码的矩阵在尾部会越来越稀疏，一方面造成资源浪费，另一方面这种表示方法直观上就很不利于模型的学习，因为它与我们实际场景是有很大的矛盾的。而RoPE虽然具有相对位置编码的性质，但是从代码部分可以看出，在构造的时候，其也是受到了最大长度的限制的。关于这一点，我无法严谨得说明，只是一点个人的想法。）。

而RoPE的巧妙之处在于，它既保留了绝对位置编码中的绝对位置信息，又保留了在内积运算下，对位置信息的相对性。

RoPE主要借助了复数的思想。为了引入复数，首先假设了在加入位置信息之前，原有的编码向量是二维行向量$q_m$和$k_n$，其中$m$和$n$是绝对位置，现在需要构造一个变换，将$m$和$n$引入到$q_m$和$k_n$中，即寻找变换：

$$\widetilde{q_m}=f(q,m),\widetilde{k_n}=f(k,n)$$
考虑到Attention的核心计算是内积：
$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$

所以，寻求的这个$f(\ast )$变换，应该具有特性：$⟨f(q,m),f(k,n)⟩=g(q,k,m-n)$

这里直接说结论，寻求的变换就是$q_m{e}^{im\theta }$，也就是给$q_m$乘以$e^{im\theta }$，相应地，$k_n$乘以$e^{in\theta }$。

具体的求解过程，请参考苏剑林老师的博客。

做了这样一个变换之后，根据复数的特性，有：

⟨ q m , k n ⟩ = R e [ q m k

快速上手Meta大语言模型LLaMA，构建“本地化ChatGPT”

近期，Meta发布了人工智能大语言模型LLaMA，包含70亿、130亿、330亿和650亿这4种参数规模的模型。其中，最小的LLaMA 7B也经过了超1万亿个tokens的训练。

本文我们将以7B模型为例，分享LLaMA的使用方法及其效果。

1. LLaMA的上手指南

这一部分，我们将step by step，展示LLaMA的使用流程。

1）下载LLaMA模型文件

以7B模型文件为例，包含内容如下：

2）克隆LLaMA项目及环境配置

git clone https://github.com/facebookresearch/llama.git
cd llama
pip install -r requirements.txt
pip install -e

如下示例中，相关操作均可通过IDP终端进行。

3) LLaMA初体验

在IDP的cell中运行如下代码，即可和LLaMA对话，接收LLaMA的回答啦！

对于7B模型：

TARGET_FOLDER=..
torchrun --nproc_per_node 1 example.py --ckpt_dir $TARGET_FOLDER/7B --tokenizer_path $TARGET_FOLDER/tokenizer.model

对于13B模型：

TARGET_FOLDER=..
torchrun --nproc_per_node 2 example.py --ckpt_dir $TARGET_FOLDER/13B --tokenizer_path $TARGET_FOLDER/tokenizer.model

对于不同大小的LLaMA模型文件，上述命令需要调整的为TARGET_FOLDER中模型文件的名称和node参数。

让我们来进行两个简单的问答测试。

Prompt 1:

Building a website can be done in 10 simple steps

LLaMA 7B模型的回答如下：

Prompt 2:

Please write a beautiful love poem

LLaMA 7B的模型回答如下：

LLaMA对这个提示词问题，自行增加了一些场景预设，如“I need some new poems for my birthday（为庆生，我需要一些新诗）"。

输入或调整提示词prompt，可在example.py文件中的prompts部分进行调整，如下图所示。

关于这两个对话示例，我们也同样放上ChatGPT的答案，供大家对比参考。

2. 从LLaMA到"开源版ChatGPT”？

LLaMA推出3天后，Nebuly AI的初创AI企业在LLaMA的基础上，基于RLHF（基于人类反馈的强化学习）进行训练，打造了对标ChatGPT的ChatLLaMA。

ChatLLaMA声称训练过程比ChatGPT快15倍。同时，支持DeepSpeed ZERO，可以加快微调速度。

DeepSpeed是一个开源的深度学习训练优化库，ZeRO是一种显存优化技术，用于提高大型模型的训练效率，如提高训练速度，降低成本和提高模型可用性等。

但ChatLLaMA并不提供LLaMA的模型权重，根据其license，也不可以商用。

---

我们在IDP开源的AI IDE中，增加了LLaMA和Stable Diffusion的示例文件，小伙伴们可以快速开启LLaMA体验之旅。

项目地址：https://github.com/BaihaiAI/IDP