文末新课《大模型项目开发线下营》秒杀！！！

第一部分 LLaMA的代码级解读：RMSNorm/SwiGLU/RoPE/Transformer

1.1 Meta发布LLaMA((7B 13B 33B 65B)：参数少但多数任务的效果好于GPT3

一直致力于LLM模型研究的国外TOP 3大厂除了OpenAI、Google，便是Meta(原来的Facebook)

Meta曾第一个发布了基于LLM的聊天机器人——BlenderBot 3，但输出不够安全，很快下线；再后来，Meta发布一个专门为科学研究设计的模型Galactica，但用户期望过高，发布三天后又下线

23年2.24日，Meta通过论文《LLaMA: Open and Efficient Foundation Language Models》发布了自家的大型语言模型LLaMA(这是LLaMA的GitHub代码地址，这是解读之一)，有多个参数规模的版本(7B 13B 33B 65B)

LLaMA只使用公开的数据(总计1.4T即1,400GB的token，其中CommonCrawl的数据占比67%，C4数据占比15%，Github、Wikipedia、Books这三项数据均都各自占比4.5%，ArXiv占比2.5%，StackExchange占比2%)，论文中提到

When training a 65B-parameter model, our code processes around 380 tokens/sec/GPU on 2048 A100 GPU with 80GB of RAM.
This means that training over our dataset containing 1.4T tokens takes approximately 21 days

且试图证明小模型在足够多的的数据上训练后，也能达到甚至超过大模型的效果

  • 比如13B参数的版本在多项基准上测试的效果好于2020年的参数规模达175B的GPT-3

  • 而对于65B参数的LLaMA，则可与DeepMind的Chinchilla(70B参数)和谷歌的PaLM(540B参数)旗鼓相当

  • 且Meta还尝试使用了论文「Scaling Instruction-Finetuned Language Models」中介绍的指令微调方法，由此产生的模型LLaMA-I，在MMLU(Massive Multitask Language Understanding，大型多任务语言理解)上要优于Google的指令微调模型Flan-PaLM-cont(62B)

1.2 代码级解读：LLaMA的模型架构——RMSNorm/SwiGLU/RoPE/Transformer

1.2.1 项目环境依赖：torch、fairscale、fire、sentencepiece

此项目给出的环境依赖有4个：

  1. torch

  2. fairscale，fairscale是用来做GPU分布的，一般是当使用DDP仍然遇到超显存的问题时使用fairscale

  3. fire，fire是一个命令行工具，用或者不用他都可以

sentencepiece，sentencepiece是用于tokenizer的工具包 「 SentencePiece 实现了subword单元（例如，字节对编码(BPE)和 unigram语言模型），并可以直接从原始句子训练字词模型(subword model)，这是对SentencePiece的解读：大模型词表扩充必备工具SentencePiece 」

# 引入 sentencepiece 库的 SentencePieceProcessor 模块，用于进行分词操作from sentencepiece import SentencePieceProcessor# 引入 logging 库的 getLogger 模块，用于生成日志from logging import getLogger# 引入 typing 库的 List 模块，用于注释函数参数或返回值的类型from typing import List# 引入 os 库，提供了大量与操作系统进行交互的接口import os
# 创建一个日志记录器logger = getLogger()
# 定义一个 Tokenizer 类class Tokenizer:# 初始化函数，参数为 SentencePiece 模型的路径def __init__(self, model_path: str):# 判断指定的模型文件是否存在assert os.path.isfile(model_path), model_path# 加载 SentencePiece 模型self.sp_model = SentencePieceProcessor(model_file=model_path)# 记录日志，提示模型加载成功logger.info(f"Reloaded SentencePiece model from {model_path}")
# 获取模型的词汇量、开始标记 ID、结束标记 ID、填充标记 IDself.n_words: int = self.sp_model.vocab_size()self.bos_id: int = self.sp_model.bos_id()self.eos_id: int = self.sp_model.eos_id()self.pad_id: int = self.sp_model.pad_id()# 记录日志，显示获取的信息logger.info(f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"# 确保模型的词汇量与词片段大小一致assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
# 编码函数，将输入的字符串编码为 token id 列表def encode(self, s: str, bos: bool, eos: bool) -> List[int]:# 检查输入的是否是字符串assert type(s) is str# 使用 SentencePiece 模型将字符串编码为 token id 列表t = self.sp_model.encode(s)# 如果需要在开头添加开始标记，就将开始标记 id 添加到列表的开头if bos:t = [self.bos_id] + t# 如果需要在结尾添加结束标记，就将结束标记 id 添加到列表的结尾if eos:t = t + [self.eos_id]# 返回 token id 列表return t
# 解码函数，将 token id 列表解码为字符串def decode(self, t: List[int]) -> str:# 使用 SentencePiece 模型将 token id 列表解码为字符串return self.sp_model.decode(t)

1.2.2 RMSNorm：对每个Transformer子层的输入进行归一化

为了提高训练的稳定性，对每个transformer子层的输入进行归一化，而不是对输出进行归一化，且使用由Zhang和Sennrich(2019)提出的RMSNorm(Root Mean Square Layer Normalization)

RMS Norm是一般LayerNorm的一种变体，可以在梯度下降时令损失更加平滑，与layerNorm相比，RMS Norm的主要区别在于去掉了减去均值的部分(re-centering)，只保留方差部分(re-scaling)

为一目了然，我们看下它们各自的归一化的表达式

  • LayerNorm

在给定一个输入特征向量后，先计算 x 的均值 μ 和标准差 σ

然后进行归一化操作： 其中的是可学习的缩放参数，来调整每个特征在归一化后的尺度或权重，最终作用是恢复归一化操作可能损失的信息，如数据的比例和分布等 而是偏移因子，可以对归一化并放缩后的数据进行偏移，使模型可以学习到一个最优的数值范围，比如在ReLU激活函数中，我们可能希望值在0以上

  • RMS Norm

首先，计算输入特征向量 a 的平方根均值 (其中，n是向量a的元素数量)

然后，对输入特征向量 a 进行归一化

此外，可选地，RMSNorm 还可以引入可学习的放缩参数 和偏移参数 ： 其代码实现为

class RMSNorm(torch.nn.Module):def __init__(self, dim: int, eps: float = 1e-6):super().__init__()// eps防止取倒数之后分母为0self.eps = epsself.weight = nn.Parameter(torch.ones(dim))

// x是输入def _norm(self, x):// torch.rsqrt是开平方并取倒数// x.pow(2)是平方/ mean(-1)是在最后一个维度(即hidden特征维度)上取平均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)// weight是末尾乘的可训练参数，即gireturn output * self.weight

至于RMS Norm为什么有用，需要求梯度进行分析，感兴趣的同学可以阅读RMS Norm的论文

1.2.3 SwiGLU替代ReLU

为了更好的理解SwiGLU，首先你得先了解什么是ReLU和GLU

  • ReLU的函数表达式为，这意味着对于所有负的输入值，ReLU函数的输出都是0，对于所有正的输入值，ReLU函数的输出等于输入值本身

  • GLU 的基本思想是引入一种称为“门”机制，该机制可以动态地控制信息的流动 这个公式意味着，对于每个输入 x，都会有一个相应的门值，这个门值由 sigmoid 函数产生，其范围在 0 到 1 之间(在正数区域接近于1，负数区域接近于0)，这个门值用于调节相应的输入值

  • 如果 接近 1，那么“门”就几乎完全开启，输入 x 的信息能够自由流动，于是 GLU 的输出接近于 x

  • 如果 接近 0，意味着“门”几乎完全关闭，即输入 x 的大部分或全部信息被阻止通过，于是 GLU 的输出接近 0

而LLaMA采用Shazeer(2020)提出的SwiGLU替换了原有的ReLU，SwiGLU的作用机制是根据输入数据的特性，通过学习到的参数自动调整信息流动的路径，具体是采用SwiGLU的Feedforward Neural Network (简称FNN，这是一种使用可学习的门控机制的前馈神经网络)

其在论文中以如下公式进行表述：

解释下这个公式

  1. 该公式先是通过Swish非线性激活函数处理 “输入和权重矩阵的乘积”

  2. 上面步骤1得到的结果和 “输入与权重矩阵的乘积” 进行逐元素的乘法 这个操作相当于在 Swish 激活的输出和第二个线性变换的输出之间引入了一个类似于GLU的“门”，这个门的值是由原始输入 通过线性变换 计算得到的，因此，它可以动态地控制 Swish 激活的输出

  3. 最后乘以权重矩阵

至于Swish激活函数可表示为

表示sigmoid函数，但其输入被缩放了倍，是一个可以学习的参数，比如下图，不同，Swish激活函数的形状则各异

打开网易新闻 查看更多图片

  • 当 趋近于 0 时，Swish 函数趋近于线性函数 y = x

  • 当 趋近于无穷大时，Swish 函数趋近于 ReLU 函数

对应论文见：Ramachandran et al., 2017

代码实现上：可以通过调用torch内置方法F.silu()实现，会在下文的FFN部分介绍

为增进大家对SwiGLU的理解，我还是举个简单的例子来说明这个过程

假设我们的输入 x 是一个二维向量 [2,3] ，权重矩阵 W 和 V 都是 2x2 矩阵，且我们简化问题，令 β =1
1. x[2,3]乘以权重矩阵 W 得到新的向量z，假设z是 [5, 4]
2. 对 xW的结果 z = [5, 4] 应用 Swish 激活函数，即Swish_1(z) = z ⊙ σ(z) = [5σ(5), 4σ(4)]
3. 然后，我们计算 xV 以得到“门”控制值 计算 xV 得到新的向量 y，假设 y = [1,0]
4. 接着，我们将 Swish_1(xW) 和 xV 做元素级别的乘法，也就是实施"门控"： (Swish_1(xW) ⊙ xV) = [5σ(5)*1, 4σ(4)*0] = [5σ(5), 0]
在这个例子中，我们可以看到 xV 的输出 [1,0] 在元素级别上控制了 Swish_1(xW) 的输出
第一个维度的门值为 1，因此 Swish_1(xW) 的第一个维度的输出能够“通过”，得到进入门控之前的结果 5σ(5)
第二个维度的门值为 0，因此 Swish_1(xW) 的第二个维度的输出被“阻止”了，结果为 0
这就是“门”的动态控制作用：它根据 xV 的输出调整 Swish_1(xW) 的输出，通过这种方式，模型可以根据输入 x 的不同，动态地调整信息流动

1.2.4 位置编码：如何彻底理解旋转位置嵌入(RoPE)

在位置编码上，删除了绝对位置嵌入，而在网络的每一层增加了苏剑林等人(2021)提出的旋转位置嵌入(RoPE)，其思想是采用绝对位置编码的形式 实现相对位置编码，且RoPE主要借助了复数的思想

先复习下复数的一些关键概念
1. 我们一般用 表示复数 ，实数a叫做复数的实部，实数b叫做复数的虚部

打开网易新闻 查看更多图片

3. 的共轭复数定义为：，也可记作，复数与其共轭的乘积等于它的模的平方，即，这是一个实数

1.2.4.1 旋转位置编码的原理

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

也就是说，我们分别为、设计操作、，使得经过该操作后，、就带有了位置、的绝对位置信息

考虑到Attention的核心计算是内积：

故我们希望的内积的结果带有相对位置信息，即寻求的这个变换，应该具有特性：

「怎么理解？很简单，当m和n表示了绝对位置之后，m与n在句子中的距离即位置差m-n，就可以表示为相对位置了，且对于复数，内积通常定义为一个复数与另一个复数的共轭的乘积」

  1. 为合理的求出该恒等式的一个尽可能简单的解，可以设定一些初始条件，比如、，然后可以先考虑二维情形，然后借助复数来求解 在复数中有，表示取实部的操作(复数 和“ 复数 的共轭即 ”之积仍是一个复数)，总之，我们需要寻找一种变换，使得

  2. 打开网易新闻 查看更多图片

  3. 那么代入方程后就得到两个方程 方程1： 方程2：Θf(q,m)−Θf(k,n) = Θg(q,k,m−n)对于方程1，代入得到(接着，再把和都设为0) 最后一个等号源于初始条件和，所以现在我们可以很简单地设，，即它不依赖于 至于方程2，同样代入得到 Θf(q,m)−Θf(k,m) = Θg(q,k,0) = Θf(q,0)−Θf(k,0)= Θ(q)−Θ(k) 这里的、是、本身的幅角，而最后一个等号同样源于初始条件 根据上式Θf(q,m)−Θf(k,m) = Θ(q)−Θ(k)，可得Θf(q,m)−Θ(q)=Θf(k,m)−Θ(k)，所以Θf(q,m)−Θ(q)的结果是一个只与m相关、跟q无关的函数，记为φ(m)，即Θf(q,m)=Θ(q)+φ(m)

  4. 接着令n=m−1代入Θf(q,m)−Θf(k,n) = Θg(q,k,m−n)，可以得到 Θf(q,m)−Θf(k,m-1) = Θg(q,k,1) 然后将 Θf(q,m) 和 Θf(k,m-1) 的等式代入Θf(q,m)=Θ(q)+φ(m)，我们可以得到 Θ(q) + φ(m) - (Θ(k) + φ(m-1)) = Θg(q,k,1)，整理一下就得到 即{φ(m)}是等差数列，设右端为θ，那么就解得φ(m)=mθ综上，我们得到二维情况下用复数表示的RoPE：

5. 所以说，寻求的变换就是，也就是给乘以，相应地，乘以 做了这样一个变换之后，根据复数的特性，有：

也就是，如果把二维向量看做复数，那么它们的内积，等于一个复数乘以另一个复数的共轭，得到的结果再取实部，代入上面的变换，也就有：

这样一来，内积的结果就只依赖于，也就是相对位置了 换言之，经过这样一番操作，通过给Embedding添加绝对位置信息，可以使得两个token的编码，经过内积变换（self-attn）之后，得到结果是受它们位置的差值，即相对位置影响的

于是，对于任意的位置为的二维向量，把它看做复数，乘以，而根据欧拉公式，有：

从而上述的相乘变换也就变成了(过程中注意：)：

打开网易新闻 查看更多图片

把上述式子写成矩阵形式：

打开网易新闻 查看更多图片

而这个变换的几何意义，就是在二维坐标系下，对向量进行了旋转，因而这种位置编码方法，被称为旋转位置编码

根据刚才的结论，结合内积的线性叠加性，可以将结论推广到高维的情形。可以理解为，每两个维度一组，进行了上述的“旋转”操作，然后再拼接在一起：

打开网易新闻 查看更多图片

由于矩阵的稀疏性，会造成计算上的浪费，所以在计算时采用逐位相乘再相加的方式进行：

打开网易新闻 查看更多图片

其中为矩阵逐位相乘操作

1.2.4.2 旋转位置编码的coding实现(分非LLaMA版和LLaMA版两种)

原理理解了，接下来可以代码实现旋转位置编码，考虑到LLaMA本身的实现不是特别好理解，所以我们先通过一份非LLaMA实现的版本，最后再看下LLaMA实现的版本

对于，非LLaMA版的实现，其核心就是实现下面这三个函数 (再次强调，本份关于RoPE的非LLaMA版的实现 与上面和之后的代码并非一体的，仅为方便理解RoPE的实现)

1.2.4.2.1 sinusoidal_position_embedding的编码实现

sinusoidal_position_embedding：这个函数用来生成正弦形状的位置编码。这种编码用来在序列中的令牌中添加关于相对或绝对位置的信息

def sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, device):# (max_len, 1)position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(-1)

# (output_dim//2)# 即公式里的i, i的范围是 [0,d/2]ids = torch.arange(0, output_dim // 2, dtype=torch.float)theta = torch.pow(10000, -2 * ids / output_dim)

# (max_len, output_dim//2)# 即公式里的：pos / (10000^(2i/d))embeddings = position * theta

# (max_len, output_dim//2, 2)embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)

# (bs, head, max_len, output_dim//2, 2)# 在bs维度重复，其他维度都是1不重复embeddings = embeddings.repeat((batch_size, nums_head, *([1] * len(embeddings.shape))))

# (bs, head, max_len, output_dim)# reshape后就是：偶数sin, 奇数cos了embeddings = torch.reshape(embeddings, (batch_size, nums_head, max_len, output_dim))embeddings = embeddings.to(device)return embeddings

一般的文章可能解释道这个程度基本就over了，但为了让初学者一目了然计，我还是再通过一个完整的示例，来一步步说明上述各个步骤都是怎么逐一结算的，整个过程和之前此文里介绍过的transformer的位置编码本质上是一回事..

为方便和transformer的位置编码做对比，故这里也假定output_dim = 512

1. 首先，我们有 ids 张量，当 output_dim 为 512 时，则

ids = [0,0, 1,1, 2,2, ..., 254,254, 255,255]

然后我们有一个基数为10000的指数运算，使用了公式 torch.pow(10000, -2 * ids / output_dim)

2. 执行 embeddings = position * theta 这行代码，它会将 position 的每个元素与 theta 的相应元素相乘，前三个元素为

3. 接下来我们将对 embeddings 的每个元素应用 torch.sin 和 torch.cos 函数 对于 torch.sin(embeddings)，我们将取 embeddings 中的每个元素的正弦值： 对于 torch.cos(embeddings)，我们将取 embeddings 中的每个元素的余弦值：

最后，torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1) 将这两个新的张量沿着一个新的维度堆叠起来，得到的 embeddings如下

最终，得到如下结果

[sin(\frac{0}{10000^{\frac{0}{512}}}), cos(\frac{0}{10000^{\frac{0}{512}}}), sin(\frac{0}{10000^{\frac{2}{512}}}), cos(\frac{0}{10000^{\frac{2}{512}}}), ..., cos(\frac{0}{10000^{\frac{510}{512}}})],[sin(\frac{1}{10000^{\frac{0}{512}}}), cos(\frac{1}{10000^{\frac{0}{512}}}), sin(\frac{1}{10000^{\frac{2}{512}}}), cos(\frac{1}{10000^{\frac{2}{512}}}), ..., cos(\frac{1}{10000^{\frac{510}{512}}})],[sin(\frac{2}{10000^{\frac{0}{512}}}), cos(\frac{2}{10000^{\frac{0}{512}}}), sin(\frac{2}{10000^{\frac{2}{512}}}), cos(\frac{2}{10000^{\frac{2}{512}}}), ..., cos(\frac{2}{10000^{\frac{510}{512}}})]

1.4.2.1.2 RoPE的编码实现

RoPE：这个函数将相对位置编码（RoPE）应用到注意力机制中的查询和键上。这样，模型就可以根据相对位置关注不同的位置

import torchimport torch.nn as nnimport torch.nn.functional as Fimport math



def RoPE(q, k):# q,k: (bs, head, max_len, output_dim)batch_size = q.shape[0]nums_head = q.shape[1]max_len = q.shape[2]output_dim = q.shape[-1]

# (bs, head, max_len, output_dim)pos_emb = sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, q.device)



# cos_pos,sin_pos: (bs, head, max_len, output_dim)# 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)cos_pos = pos_emb[..., 1::2].repeat_interleave(2, dim=-1) # 将奇数列信息抽取出来也就是cos 拿出来并复制sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1) # 将偶数列信息抽取出来也就是sin 拿出来并复制

# q,k: (bs, head, max_len, output_dim)q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)q2 = q2.reshape(q.shape) # reshape后就是正负交替了

# 更新qw, *对应位置相乘q = q * cos_pos + q2 * sin_pos

k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)k2 = k2.reshape(k.shape)# 更新kw, *对应位置相乘k = k * cos_pos + k2 * sin_pos

return q, k

老规矩，为一目了然起见，还是一步一步通过一个示例来加深理解

1. sinusoidal_position_embedding函数生成位置嵌入。在output_dim=512的情况下，每个位置的嵌入会有512个维度，但为了简单起见，我们只考虑前8个维度，前4个维度为sin编码，后4个维度为cos编码。所以，我们可能得到类似以下的位置嵌入

# 注意，这只是一个简化的例子，真实的位置嵌入的值会有所不同。pos_emb = torch.tensor([[[[0.0000, 0.8415, 0.9093, 0.1411, 1.0000, 0.5403, -0.4161, -0.9900],[0.8415, 0.5403, 0.1411, -0.7568, 0.5403, -0.8415, -0.9900, -0.6536],[0.9093, -0.4161, -0.8415, -0.9589, -0.4161, -0.9093, -0.6536, 0.2836]]]])

2. 然后，我们提取出所有的sin位置编码和cos位置编码，并在最后一个维度上每个位置编码进行复制

sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1) # 提取出所有sin编码，并在最后一个维度上复制cos_pos = pos_emb[..., 1::2].repeat_interleave(2, dim=-1) # 提取出所有cos编码，并在最后一个维度上复制

3. 更新query向量 我们首先构建一个新的q2向量，这个向量是由原来向量的负的cos部分和sin部分交替拼接而成的 我们用cos_pos对q进行元素级乘法，用sin_pos对q2进行元素级乘法，并将两者相加得到新的query向量

q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1).flatten(start_dim=-2)# q2: tensor([[[[-0.2, 0.1, -0.4, 0.3, -0.6, 0.5, -0.8, 0.7],# [-1.0, 0.9, -1.2, 1.1, -1.4, 1.3, -1.6, 1.5],# [-1.8, 1.7, -2.0, 1.9, -2.2, 2.1, -2.4, 2.3]]]])

q = q * cos_pos + q2 * sin_pos

公式表示如下

打开网易新闻 查看更多图片

4. 更新key向量 对于key向量，我们的处理方法与query向量类似

k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1).flatten(start_dim=-2)# k2: tensor([[[[-0.15, 0.05, -0.35, 0.25, -0.55, 0.45, -0.75, 0.65

1.4.2.1.3 attention的编码实现

attention：这是注意力机制的主要功能

  • 首先，如果use_RoPE被设置为True，它会应用RoPE，通过取查询和键的点积（并进行缩放）

  • 然后，进行softmax操作来计算注意力分数，以得到概率，输出是值的加权和，权重是计算出的概率

  • 最后，旋转后的q和k计算点积注意力后，自然就具备了相对位置信息


def attention(q, k, v, mask=None, dropout=None, use_RoPE=True):# q.shape: (bs, head, seq_len, dk)# k.shape: (bs, head, seq_len, dk)# v.shape: (bs, head, seq_len, dk)

if use_RoPE:# 使用RoPE进行位置编码q, k = RoPE(q, k)

d_k = k.size()[-1]

# 计算注意力权重# (bs, head, seq_len, seq_len)att_logits = torch.matmul(q, k.transpose(-2, -1))att_logits /= math.sqrt(d_k)

if mask is not None:# 对权重进行mask，将为0的部分设为负无穷大att_scores = att_logits.masked_fill(mask == 0, -1e-9)

# 对权重进行softmax归一化# (bs, head, seq_len, seq_len)att_scores = F.softmax(att_logits, dim=-1)

if dropout is not None:# 对权重进行dropoutatt_scores = dropout(att_scores)

# 注意力权重与值的加权求和# (bs, head, seq_len, seq_len) * (bs, head, seq_len, dk) = (bs, head, seq_len, dk)return torch.matmul(att_scores, v), att_scores



if __name__ == '__main__':# (bs, head, seq_len, dk)q = torch.randn((8, 12, 10, 32))k = torch.randn((8, 12, 10, 32))v = torch.randn((8, 12, 10, 32))

# 进行注意力计算res, att_scores = attention(q, k, v, mask=None, dropout=None, use_RoPE=True)

# 输出结果的形状# (bs, head, seq_len, dk), (bs, head, seq_len, seq_len)print(res.shape, att_scores.shape)






接下来，我们再来看下LLaMA里是怎么实现这个旋转位置编码的，具体而言，LLaMA 的model.py文件里面实现了旋转位置编码(为方便大家理解，我给相关代码 加了下注释)首先，逐一实现这三个函数precompute_freqs_cisreshape_for_broadcastapply_rotary_emb

# 预计算频率和复数的函数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_cis = torch.polar(torch.ones_like(freqs), freqs) # 计算复数return freqs_cis # 返回复数


# 重塑的函数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_cis: torch.Tensor,) -> Tuple[torch.Tensor, torch.Tensor]:xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # 将xq视为复数xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # 将xk视为复数freqs_cis = reshape_for_broadcast(freqs_cis, xq_) # 重塑复数的形状xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # 计算xq的输出xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) # 计算xk的输出return xq_out.type_as(xq), xk_out.type_as(xk) # 返回xq和xk的输出
之后，在注意力机制的前向传播函数中调用上面实现的第三个函数 apply_rotary_emb，赋上位置信息 (详见下文1.2.5节)

# 对Query和Key应用旋转嵌入xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

1.2.5 Transform架构的实现：Attention计算、SA、FFN

LLaMA和GPT一样，都是基于Transformer这个架构，通常，我们在构建transformer时，是按Block构建的，每个transformer Block包含SA和FFN两部分，然后再通过堆叠block的形式，构建起整个transformer网络，LLaMA也是这样做的

回顾一下Attention计算的总体过程是：

  1. 输入，分别经过三个Linear得到

  2. 在 和中加入旋转位置编码

  3. 缓存 和

  4. 计算

其中有一个细节就是缓存机制，它设计的目的是在generate时减少token的重复计算。简单解释一下，就是在计算第n个token特征的时候，需要用到第个token，即每次生成时，需要知道前面所有的过往信息，如果每次都从头算的话，那就会造成极大的浪费，所以就没算一个位置的信息，就把它缓存下来

接下来，我们来看下代码实现，首先是SA(self-attention)部分：

class Attention(nn.Module):def __init__(self, args: ModelArgs):super().__init__()
# 设置本地注意力头的数量self.n_local_heads = args.n_heads // fs_init.get_model_parallel_world_size()# 每个注意力头的维度self.head_dim = args.dim // args.n_heads

# Query投影层self.wq = ColumnParallelLinear(args.dim,args.n_heads * self.head_dim,bias=False,gather_output=False,init_method=lambda x: x,# Key投影层self.wk = ColumnParallelLinear(args.dim,args.n_heads * self.head_dim,bias=False,gather_output=False,init_method=lambda x: x,# Value投影层self.wv = ColumnParallelLinear(args.dim,args.n_heads * self.head_dim,bias=False,gather_output=False,init_method=lambda x: x,# 输出投影层self.wo = RowParallelLinear(args.n_heads * self.head_dim,args.dim,bias=False,input_is_parallel=True,init_method=lambda x: x,

# 使用零初始化键缓存self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)).cuda()# 使用零初始化值缓存self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)).cuda()

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

# 将形状调整为[bsz, seqlen, n_local_heads, head_dim]xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)

# 对Query和Key应用旋转嵌入xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

# 将缓存键和值转换为xq的设备类型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] = xkself.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]

# 转置xq、键和值的维度xq = xq.transpose(1, 2)keys = keys.transpose(1, 2)values = values.transpose(1, 2)

# 计算注意力分数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, slen, cache_len + slen)scores = F.softmax(scores.float(), dim=-1).type_as(xq)

# 使用注意力分数加权求和得到输出output = torch.matmul(scores, values) # (bs, n_local_heads, slen, head_dim)output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

# 应用输出投影return self.wo(output)

然后是前馈网络FFN部分，需要注意的点就是采用的激活函数，以及激活函数的位置

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

class FeedForward(nn.Module):def __init__(self,dim: int,hidden_dim: int,multiple_of: int,super().__init__()

# 初始化隐藏层的维度为输入维度的2/3hidden_dim = int(2 * hidden_dim / 3)# 调整隐藏层维度为multiple_of的倍数hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

# 第一个线性层self.w1 = ColumnParallelLinear(dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x

# 第二个线性层self.w2 = RowParallelLinear(hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x

# 第三个线性层self.w3 = 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))

这里与常见模型中的FFN做一下简单的对比

  • BART中的FFN，用的是fc->act->fc，用了两层全连接

  • GPT中的FFN，用的是conv1D->act->conv1D，也是只用了两层

  • 而LLaMA中的FFN采用了三个全连接层以实现FFNSwiGLU，即

然后将SA和FFN这两部分拼在一起就是一个transformer block

import torchimport torch.nn as nnfrom typing import Optional

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=4 * args.dim, multiple_of=args.multiple_of) # 前馈神经网络模块self.layer_id = layer_id # 当前层的IDself.attention_norm = RMSNorm(args.dim, eps=args.norm_eps) # 注意力模块的归一化self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) # 前馈神经网络模块的归一化
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):# 输入x经过self-attention之后，做Add&Normh = x + self.attention.forward(self.attention_norm(x), start_pos, freqs_cis, mask)# 上一步的输出h作为输入，经过前馈神经网络Feed forward之后，做Add&Normout = h + self.feed_forward.forward(self.ffn_norm(h))return out

最后利用torch的module list将transformer block进行堆叠，拼上最前头的embedding部分，就是一个完整的transformer decoder结构了

import torchimport torch.nn as nnfrom typing import Optional

class Transformer(nn.Module):def __init__(self, params: ModelArgs):super().__init__()

# 初始化参数self.params = paramsself.vocab_size = params.vocab_size # 词汇表大小self.n_layers = params.n_layers # Transformer模型的层数

# 词嵌入层self.tok_embeddings = ParallelEmbedding(params.vocab_size, params.dim, init_method=lambda x: x

# Transformer的各个层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 = ColumnParallelLinear(params.dim, params.vocab_size, bias=False, init_method=lambda x: x

# 预计算的频率矩阵self.freqs_cis = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len * 2

@torch.inference_mode()def forward(self, tokens: torch.Tensor, start_pos: int):_bsz, seqlen = tokens.shape

# Token嵌入和位置编码h = self.tok_embeddings(tokens)self.freqs_cis = self.freqs_cis.to(h.device)freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

# 生成上三角的mask矩阵（为decoder模型防止标签泄漏）mask = Noneif seqlen > 1:mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=tokens.device)mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)

# 逐层计算Transformerfor layer in self.layers:h = layer(h, start_pos, freqs_cis, mask)h = self.norm(h)output = self.output(h[:, -1, :]) # 只计算最后一个位置的logitsreturn output.float()

接着看下生成过程，如下：

  1. 对prompts进行tokenize，得到token ids；

  2. 计算当前batch的最大长度total_len，用来创建输入的token tensor，最大长度不能超过前文所述缓存的大小；

  3. 从当前batch中，最短的一个prompt的位置，作为生成的开始位置，开始生成；

  4. 输入的token tensor传入transformer模型，计算logits，得到形状为(batch_size, hidden_size)的logits（transformer最后一层的输出）；

  5. softmax+top_p采样，得到当前预测的token，并更新当前位置，准备预测下一个token；

  6. 解码得到生成的文本

代码如下

class LLaMA:def __init__(self, model: Transformer, tokenizer: Tokenizer):self.model = modelself.tokenizer = tokenizer
def generate(self,prompts: List[str],max_gen_len: int,temperature: float = 0.8,top_p: float = 0.95,) -> List[str]:# 获取批处理大小bsz = len(prompts)# 获取模型参数params = self.model.params# 检查批处理大小是否在允许的最大批处理大小范围内assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)
# 使用分词器对提示进行编码为标记prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
# 查找提示标记的最小和最大大小min_prompt_size = min([len(t) for t in prompt_tokens])max_prompt_size = max([len(t) for t in prompt_tokens])
# 计算要生成的标记的总长度total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)
# 创建一个张量来存储生成的标记，填充为填充标记tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cuda().long()# 将提示标记复制到标记张量中for k, t in enumerate(prompt_tokens):tokens[k, : len(t)] = torch.tensor(t).long()# 创建一个掩码以识别输入文本input_text_mask = tokens != self.tokenizer.pad_id# 设置生成的起始位置start_pos = min_prompt_sizeprev_pos = 0# 逐个生成标记for cur_pos in range(start_pos, total_len):# 通过模型进行前向传递以获取logitslogits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)if temperature > 0:# 对logits应用温度并计算概率probs = torch.softmax(logits / temperature, dim=-1)# 使用top-p采样抽样下一个标记next_token = sample_top_p(probs, top_p)else:# 选择概率最高的标记next_token = torch.argmax(logits, dim=-1)next_token = next_token.reshape(-1)# 只有在已经生成了提示的情况下才替换标记next_token = torch.where(input_text_mask[:, cur_pos], tokens[:, cur_pos], next_tokentokens[:, cur_pos] = next_tokenprev_pos = cur_pos
# 将生成的标记解码为文本decoded = []for i, t in enumerate(tokens.tolist()):# 将标记截断到最大生成长度t = t[: len(prompt_tokens[i]) + max_gen_len]# 将标记截断到如果存在结束标记try:t = t[: t.index(self.tokenizer.eos_id)]except ValueError:pass# 将标记解码为文本decoded.append(self.tokenizer.decode(t))return decoded
def sample_top_p(probs, p):# 按降序对概率进行排序probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)# 计算概率的累积和probs_sum = torch.cumsum(probs_sort, dim=-1)# 创建一个掩码以过滤累积概率超过p的标记mask = probs_sum - probs_sort > p# 将被过滤的标记的概率设置为0probs_sort[mask] = 0.0# 归一化概率probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))# 使用修改后的概率进行抽样下一个标记next_token = torch.multinomial(probs_sort, num_samples=1)# 收集抽样标记的原始索引next_token = torch.gather(probs_idx, -1, next_token)return next_token

1.3 LLaMA的Optimizer设计、模型加速优化与微型版本

在Optimizer设计上

  • 该模型使用AdamW优化器(Loshchilov和Hutter，2017)进行训练，超参数设置为β1=0.9，β2=0.95 此外，使用余弦学习率方式，使最终学习率等于最大学习率的10%，以及使用0.1的权重衰减和1.0的梯度剪裁，和2000个warm up策略，使得可以根据模型的大小改变学习率和批次大小

在模型的加速优化方面

  1. 首先，使用一个高效的因果多头注意力方式的实现，灵感来自Rabe和Staats(2021)以及Dao等人(2022)，这个实现可在xformers库中找到，可以有效减少内存的使用和计算 具体原理为通过不存储注意力权重和不计算由于语言建模任务的因果性质而被掩盖的键/查询分数来实现的

  2. 其次，为了进一步提高训练效率，减少了在check point的后向传递中重新计算的激活量，在实现上，通过手动实现trasnformer层的后向函数来进行操作 为了充分受益于这种优化，还通过如Korthikanti等人(2022)中采用的方法，进行使用模型和序列并行来减少模型的内存使用

  3. 最后，该工作还尽可能地重叠激活的计算和GPU之间在网络上的通信 最终的优化性能效果为：当训练一个65B参数的模型时，代码在2048A100的GPU上处理大约380个token/秒/GPU，并耗费80GB的内存，这意味着对包含1.4Ttoken的数据集进行训练大约花费了21天

LLaMA发布不久后，一些研究者基于它做了不少工作

  • 一开始最小参数7B的模型也需要近30GB的GPU才能运行，但通过比特和字节库进行浮点优化，能够让模型在单个NVIDIA RTX 3060(显存一般12G)上运行

  • 之后，GitHub 上的一名研究人员甚至能够在Ryzen 7900X CPU上运行LLM的7B 版本，每秒能推断出几个单词

  • 再之后，有研究者推出了llama.cpp，无需 GPU，就能运行 LLaMA llama.cpp 项目实现了在MacBook上运行 LLaMA，还有开发者成功的在 4GB RAM 的树莓派上运行了 LLaMA 7B

第二部分 限于篇幅此处省略，想要私七月在线位老师微信（或找苏苏老师：julyedukefu008）

第三部分 更强的LLaMA 2开源，可直接商用

3.1 LLaMA2简介：相比LLaMA1代——1.4倍token，2倍上下文

LLAMA 2 (项目地址、论文地址、 )，是 LLAMA 1 的更新版本

  • 模型结构采用了 Llama 1 的大部分预训练设置和模型架构，比如使用标准Transformer 架构，使用 RMSNorm 应用预归一化、使用 SwiGLU 激活函数和旋转位置嵌入RoPE

  • 训练数据使用一种新的混合的公开可用数据进行训练，训练数据规模是2T个token，相比1代的1.4T多出了40%

  • 上下文长度上下文长度达到了4096，相比1代的2048直接翻了一倍

  • 模型种类目前 LLAMA 2 的系列模型有 7B、13B、70B 三种(34B的后续发布) 值得特别注意的是，其中的70B模型采用了分组查询注意力(grouped-query attention，简称GQA)「Transformer原始论文中用的多头注意力(MHA)、ChatGLM2-6B则用的多查询注意力(Multi-query attention，简称MQA)」

打开网易新闻 查看更多图片

同时 Meta 还发布了 LLaMA 2-CHAT，其是基于 LLAMA 2 针对对话场景微调的版本，同样 7B、13B 和 70B 参数三个版本，具体的训练方法与ChatGPT类似

打开网易新闻 查看更多图片

  1. 先是监督微调LLaMA2得到SFT版本 (接受了成千上万个人类标注数据的训练，本质是问题-答案对 )

  2. 然后使用人类反馈强化学习(RLHF)进行迭代优化 先训练一个奖励模型 然后在奖励模型/优势函数的指引下，通过拒绝抽样(rejection sampling)和近端策略优化(PPO)的方法迭代模型的生成策略

LLAMA 2 的性能表现更加接近 GPT-3.5，Meta 也承认距离 GPT-4 和 PaLM 2 等领先非开源模型还有差距

打开网易新闻 查看更多图片

Meta 在技术报告中详细列出了 LLAMA 2 的性能、测评数据，以及分享了重要的训练方法，具体详见原论文

打开网易新闻 查看更多图片

3.2 LLaMA2之分组查询注意力——Grouped-Query Attention

自回归解码的标准做法是缓存序列中先前标记的键 (K) 和值 (V) 对，从而加快注意力计算速度

然而，随着上下文窗口或批量大小的增加，多头注意力 (MHA)模型中与 KV 缓存大小相关的内存成本显着增长

对于较大的模型，KV 缓存大小成为瓶颈，键和值投影可以在多个头之间共享，而不会大幅降低性能，可以使用

  • 具有单个 KV 投影的原始多查询格式(MQA)ChatGLM2-6B即用的这个，详见此文《中文模型的奋起直追：MOSS、baichuan-7B/13B和ChatGLM2-6B/13B的原理、部署与微调》的3.1节 不过，多查询注意(Multi-query attention，简称MQA)只使用一个键值头，虽大大加快了解码器推断的速度，但MQA可能导致质量下降，而且仅仅为了更快的推理而训练一个单独的模型可能是不可取的

  • 或具有多个 KV 投影的分组查询注意力(grouped-query attention，简称GQA)，速度快 质量高 23年，还是Google的研究者们提出了一种新的方法，即分组查询注意(GQA，论文地址为：GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints)，这是一种多查询注意的泛化，它通过折中(多于一个且少于查询头的数量，比如4个)键值头的数量，使得经过强化训练的GQA以与MQA相当的速度达到接近多头注意力的质量

打开网易新闻 查看更多图片

经实验论证，GQA 变体在大多数评估任务上的表现与 MHA 基线相当，并且平均优于 MQA 变体

3.3 Llama 2-Chat中的RLHF：依然是三阶段训练方式

3.3.1 监督微调(SFT)

在SFT的数据上

  1. 他们先是重点收集了几千个高质量 SFT 数据示例 (注意：很多新闻稿会说SFT的数据达到百万以上，这就是没仔细看论文的结果，论文之意是胜过百万低质量的数据，As a result, we focused first on collecting several thousand examples of high-quality SFT data. By setting aside millions of examples from third-party datasets and using fewer buthigher-quality examples from our own vendor-based annotation efforts, our results notably improved)

  2. 之后发现几万次的SFT标注就足以获得高质量的结果，最终总共收集了27540条用于SFT的标注数据 (We found that SFT annotations in the order of tens ofthousands was enough to achieve a high-quality result. We stopped annotating SFT after collecting a total of 27,540 annotations.)

在微调过程中

  • 每个样本都包括一个prompt和一个response(说白了，就是问题-答案对，和instructGPT/ChatGPT本身的监督微调是一个本质)，且为确保模型序列长度得到正确填充，Meta 将训练集中的所有prompt和response连接起来。他们使用一个特殊的 token 来分隔prompt和response片段，利用自回归目标，将来自用户提示的 token 损失归零，因此只对答案 token 进行反向传播，最后对模型进行了 2 次微调

  • 微调过程中的参数则如此设置：we use a cosine learning rate schedule with an initiallearning rate of 2 ×10−5 , a weight decay of 0.1, a batch size of 64, and a sequence length of 4096 token

打开网易新闻 查看更多图片

3.3.2 训练两个奖励模型：一个偏实用 一个偏安全

下表 6 报告了 Meta 长期以来收集到的奖励建模数据的统计结果，并将其与多个开源偏好数据集进行了对比。他们收集了超过 100 万个基于人类应用指定准则的二元比较的大型数据集，也就是奖励建模数据

打开网易新闻 查看更多图片

关于奖励数据

  1. prompt和response中的标记数因文本领域而异，比如摘要和在线论坛数据的prompt通常较长，而对话式的prompt通常较短。与现有的开源数据集相比，本文的偏好数据具有更多的对话回合，平均长度也更长

  2. 奖励模型将模型响应及其相应的提示(包括前一轮的上下文)作为输入，并输出一个标量分数来表示模型生成的质量(例如有用性和安全性)，利用这种作为奖励的响应得分，Meta 在 RLHF 期间优化了 Llama 2-Chat，以更好地与人类偏好保持一致，并提高有用性和安全性 在每一批用于奖励建模的人类偏好标注中，Meta 都拿出 1000 个样本作为测试集来评估模型，并将相应测试集的所有prompt的集合分别称为实用性和安全性 (很多新闻稿会翻译成元实用、元安全，其实没必要加个“元”字，你理解为是Meta内部定义的“实用”与“安全”两个概念即可)

故为了兼顾和平衡模型的实用性和安全性，LLaMA 2团队训练了两个独立的奖励模型

一个针对实用性(称为实用性RM)进行了优化，在内部所有偏实用的奖励数据集上进行训练，并结合从内部偏安全的奖励数据集和开源安全性数据集中统一采样的同等部分剩余数据theHelpfulness reward model is eventually trained on all Meta Helpfulness data, combined with an equalparts of the remaining data uniformly sampled from Meta Safety and from the open-source datasets

另一个针对安全性(安全性RM)进行了优化，在内部所有偏安全的奖励数据和人类无害数据上进行训练，并以90/10的比例混合内部偏实用的奖励数据和开源实用性数据TheMeta Safety reward model is trained on all Meta Safety and Anthropic Harmless data, mixed with MetaHelpfulness and open-source helpfulness data in a 90/10 proportion.

We found that the setting with 10% helpfulness data is especially beneficial for the accuracy on samples where both the chosen and rejectedresponses were deemed safe

并通过预训练的LLaMA 2初始化奖励模型(意味着奖励模型的架构与参数与预训练模型一致，只是用于下一个token预测的分类头被替换为用于输出标量奖励的回归头)，因为它确保了两个模型都能从预...