【大模型】LLaMa系列演进及源码解析

本文转载自： LLaMa-1/2/3 原理+源码——拆解 (KV-Cache, RoPE, RMSNorm, GQA, SwiGLU)

一、背景介绍

主流的大语言模型都采用了Transformer架构，它是一个基于多层Self-attention的神经网络模型。

原始的Transformer由编码器（Encoder）和解码器（Decoder）两个部分构成，同时，这两个部分也可以独立使用。

Llama模型与GPT-2类似，也是采用了基于Decoder-Only的架构。在原始Vanilla Transformer Decoder的基础上，Llama进行了如下改动：

• 为了增强训练稳定性，前置了层归一化（Pre-normalization），并使用RMSNorm 作为层归一化方法。
• 为了提高模型性能，采用SwiGLU 作为激活函数。
• 为了更好地建模长序列数据，采用RoPE 作为位置编码。
• 为了平衡效率和性能，部分模型采用了GQA分组查询注意力机制（Grouped-Query Attention, GQA）。
• 并且将self-attention改进为使用KV-Cache的Grouped Query。
• 每个版本的Llama由于其隐藏层的大小、层数的不同，均有不同的变体。

二、LLaMa-1/2/3 演进

接下来，我们将展开看下每个版本的不同变体：

2.1 LLaMa-1系列的演进

• LLaMa-1：Meta开源的Pre-trained Model，模型参数从 7B、13B、32B、65B 不等，LLaMa-7B在大多数基准测试上超过了Text-davinci-003（即GPT3-173B），相比于ChatGPT或者GPT4来说，LLaMa可能效果上还有差距，目前hugging face已集成了LLaMa的代码实现和开源模型。学术界和工业界都可以在此基础上进行学习和研究。
  

• Alpaca：斯坦福在LLaMa-7B的基础上监督微调出来的模型，斯坦福是用OpenAI的Text-davinci-003（即GPT3-173B）的API配合self-instruct技术，使用175个提示语种子自动生成了52K条提示-回复的指示数据集，在LLaMa-7B上微调得到的模型，在8张80G的A100上训练了3小时。

• Vicuna：在LLaMa-13B的基础上使用监督微调得到的模型，数据集来自于ShareGPT 产生的用户对话数据，共70K条。使用Pytorch FSDP在8张A100上训练了一天。相较于Alpaca，Vicuna在训练中将序列长度由512扩展到了2048，并且通过梯度检测和flash attention来解决内存问题；调整训练损失考虑多轮对话，并仅根据模型的输出进行微调。通过GPT4来打分评测，Vicuna可以达到ChatGPT 90%的效果。

2.2 LLaMa2系列的演进

• LLaMa2：采用了Llama 1的大部分预训练设置和模型架构，有 7B、13B、34B、70B四个参数量版本。LLaMa2和LLaMa1的最大差别是：Llama-2将预训练的语料扩充到了 2T token语料，同时将模型的上下文长度从2,048翻倍到了4,096，并在训练34B、70B的模型中引入了分组查询注意力机制（grouped-query attention, GQA）。

• LLaMa-2 Chat：有了更强大的基座模型Llama-2，Meta通过进一步的有监督微调（Supervised Fine-Tuning, SFT）、基于人类反馈的强化学习（Reinforcement Learning with Human Feedback, RLHF）等技术对模型进行迭代优化（Pertrain -> SFT -> RLHF），并发布了面向对话应用的微调系列模型 Llama-2 Chat。

• Code-LLaMa：得益于Llama-2的优异性能，Meta在2023年8月发布了专注于代码生成的Code-Llama，共有7B、13B、34B和70B四个参数量版本。

2.3 LLaMa3系列的演进

• LLaMa3：包括8B和70B两个参数量版本。除此之外，Meta还透露，400B的Llama-3还在训练中。相比Llama-2，Llama-3支持8K长文本，并采用了一个编码效率更高的tokenizer，词表大小为128K。在预训练数据方面，Llama-3使用了超过15T token的语料，这比Llama 2的7倍还多。
  – 小型模型具有8B参数，其性能略优于Mistral 7B和Gemma 7B；
  – 中型模型则拥有70B参数，其性能介于ChatGPT 3.5和GPT 4之间；
  – 大型模型规模达到400B，目前仍在训练中，旨在成为一个多模态、多语言版本的模型，预期性能应与GPT 4或GPT 4V相当。

三、LLaMa关键技术点

3.1 Embedding

Embedding的过程：word -> token_id -> embedding_vector

其中第一步转化使用tokenizer的词表进行，第二步转化使用 learnable 的 Embedding layer。

3.2 均方根层归一化 (RMS Norm)

我们在之前的任务中，使用的归一化方式主要是 Batch Norm（批归一化） 和 Layer Norm（层归一化） ，那么LLama中的一个改进点在于使用 RMS Norm（均方根层归一化），下面我们看看三者的区别：

• Layer Norm：层归一化(LayerNorm)对Transformer等模型来说非常重要，它可以帮助稳定训练并提升模型收敛性。LayerNorm针对一个样本所有特征计算均值和方差，然后使用这些来对样本进行归一化：
  
  这里 $x=(x_1,x_2,\cdots ,x_H)$ 表示某个时间步LN层的输入向量表示，向量维度为$H$；$h$ 是LN层的输出； $g$ , $b$ 是两个可学习的参数。

为什么层归一化有用？一些解释如下：

1. 减少内部协变量偏移（Internal Covariate Shift）： 内部协变量偏移是指在深度神经网络的训练过程中，每一层输入的分布会发生变化，导致网络的训练变得困难。层归一化通过对每一层的输入进行归一化处理，可以减少内部协变量偏移，使得每一层的输入分布更加稳定。
2. 稳定化梯度： 层归一化有助于保持每一层输出的均值和方差稳定，从而使得梯度的传播更加稳定。这有助于减少梯度消失或梯度爆炸的问题，提高梯度在网络中的流动性，加快训练速度。
3. 更好的参数初始化和学习率调整： 通过层归一化，每一层的输入分布被归一化到均值为0、方差为1的标准正态分布，这有助于更好地初始化网络参数和调整学习率。参数初始化与学习率调整的稳定性对模型的训练效果至关重要。
4. 增强模型的泛化能力： 层归一化可以减少网络对训练数据分布的依赖，降低了过拟合的风险，从而提高模型的泛化能力。稳定的输入分布有助于模型更好地适应不同数据集和任务。

• Batch Norm：Batch Norm 和 Layer Norm都是一样减去均值Mean，除以方差Var，最终将归一化为正态分布N(0,1)。只不过两者是在不同的维度(batch还是feature)求均值和方差，（其中，减均值：re-centering 将均值mean变换为0，除方差：re-scaling将方差varance变换为1）。

• RMS Norm（Root MeanSquare Layer Norm）：RMS Norm（均方根层归一化）认为，虽然LayerNorm很好，但是它每次需要计算均值和方差。实际上，Layer Norm成功的原因是re-scaling，因为方差Var计算的过程中使用了均值Mean。因此，RMS Norm不再使用均值Mean，而是构造了一个特殊的统计量RMS代替方差Var。
  
  这里，$gi$ 为可学习的缩放因子，偏移参数$bi$移除掉了。
  单看上式的话，相当于仅使用 $x$ 的均方根来对输入进行归一化，它简化了层归一化的计算，使得计算量更小，同时还有可能带来性能上的提升。

关于RMSNorm的具体介绍，可以参考博客：Llama改进之——均方根层归一化RMSNorm

RMSNorm 在HuggingFace Transformer 库中代码实现如下所示：

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps # eps 防止取倒数之后分母为0
    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        # weight 是末尾乘的可训练参数, 即g_i
        return (self.weight * hidden_states).to(input_dtype)

为了使得模型训练过程更加稳定，GPT-2 相较于GPT 就提出了将Layer Norm前置，将第一个层归一化移动到多头自注意力层之前，第二个层归一化也移动到了全连接层之前，同时残差连接的位置也调整到了多头自注意力层与全连接层之后。层归一化中也采用了RMSNorm 归一化函数。

3.3 旋转位置编码 (Rotary Positional Encodding, RoPE)

参考博客：十分钟读懂旋转编码（RoPE）

旋转位置编码（Rotary Position Embedding，RoPE）是论文Roformer: Enhanced Transformer With Rotray Position Embedding 提出的一种能够将相对位置信息依赖集成到 self-attention 中并提升 transformer 架构性能的位置编码方式。

和相对位置编码相比，RoPE 具有更好的外推性，目前是大模型相对位置编码中应用最广的方式之一。

大模型的外推性：是指大模型在训练时和预测时的输入长度不一致，导致模型的泛化能力下降的问题。例如，如果一个模型在训练时只使用了512个 token 的文本，那么在预测时如果输入超过512个 token，模型可能无法正确处理。这就限制了大模型在处理长文本或多轮对话等任务时的效果。

【绝对位置编码】：

在传统的Transformer中，主要使用的是绝对位置编码，绝对Positional Encodding的使用过程：word -> token_id -> embedding_vector + position_encodding -> Encoder_Input

其中第一步转化使用tokenizer的词表进行，第二步转化使用 learnable 的 Embedding layer。将得到的embedding_vector 和 position_encodding 进行element-wise的相加，然后才做为input送入LLM的encoder。

对比Absolute PE 和 Relative PE：

• Absolute PE 绝对位置编码：每次单独1个token的PE，每个token的PE之间没有关系，是一组固定的vector，反映每个token在sequence中的绝对位置。计算 query, key 和 value 向量之前加在输入序列X上，$Q/K/V=W_{q/k/v}(X+P)$，经典的位置编码向量 P 的计算方式是使用 Sinusoidal 函数。

• Relative PE 相对位置编码：每次处理2个token的PE，只在计算attention时使用（在query@key时加在key上），反映2个token的相关度。

【旋转位置编码（RoPE）】：
旋转位置编码（RoPE）：RoPE 借助了复数的思想，出发点是通过绝对位置编码的方式，实现token间的相对位置编码。其目标是通过下述 f 运算，计算self-attention前，来给q，k 添加，其在sequence中的绝对位置信息m和n，得到qm 和kn，然后进行qm@kn：

公式比较复杂，总结来说，RoPE 的 self-attention 操作的流程是：对于 token 序列中的每个词嵌入向量，首先计算其对应的 query 和 key 向量，然后对每个 token 位置都计算对应的旋转位置编码，接着对每个 token 位置的 query 和 key 向量的元素按照 两两一组 应用旋转变换，最后再计算 query 和 key 之间的内积得到 self-attention 的计算结果。

论文中有个很直观的图片展示了旋转变换的过程：

3.4 SwiGLU Function

SwiGLU并不是一种全新的算法或理论，而是对现有Transformer架构中的FFN层的一种改进。在Transformer中，FFN是实现前馈传播的关键部分，通过两层全连接层和ReLU激活函数，实现从输入到输出的映射。

Transformer中的FFN层定义如下：

然而，SwiGLU对这一结构进行了优化，将第一层全连接和ReLU激活函数替换为两个权重矩阵和输入的变换，再配合Swish激活函数进行哈达马积操作。

SwiGLU的公式如下：

其中，σ(x) 是Sigmoid 函数。下图给出了Swish 激活函数在参数β 不同取值下的形状。可以看到当β 趋近于0 时，Swish 函数趋近于线性函数y = x，当β 趋近于无穷大时，Swish 函数趋近于ReLU 函数，β 取值为1 时，Swish 函数是光滑且非单调。

上述公式中包含有 Swish 部分和 GLU 部分：

• swi指的是swish非线性激活函数：
  
• GLU指的是 gated linear unit，输入的向量x分别经过两个linear层，其中一个需要经过非线性激活函数，然后将两者对应元素相乘。
  
  SwiGLU实际上就是讲 Swish和GLU组合起来。

3.5 KV-Cache

参考博客：Transformer系列：图文详解KV-Cache，解码器推理加速优化

KV-Cache是一种加速Transformer推理的策略，几乎所有自回归模型都内置了KV-Cache，理解KV-Cache有助于更深刻地认识Transformer中注意力机制的工作方式。

KV Cache，即键-值缓存，是一种用于存储键值对数据的缓存机制。在语言模型的推理过程中，经常需要多次访问相同的数据，而KV Cache通过将这些数据缓存到内存中，提供了快速的数据访问速度，从而加速推理过程。该技术仅应用于解码阶段。如 decode only 模型（如 GPT3、Llama 等）、encode-decode 模型（如 T5）的 decode 阶段，像 Bert 等非生成式模型并不适用。

【自回归推理过程】：
自回归模型采用shift-right的训练方式，用前文预测下一个字/词，并且前文中的最后一个词经过解码器的表征会映射为其下一个待预测词的概率分布。在训练阶段，句子完整输入给网络，所有位置下的token并行计算。
同理，在预测推理阶段也可以将前文prompt完整输入给训练好的模型，取最后一个位置的表征作为下一个token的概率分布，再通过采样策略确认下一个token，最终将token拼接到前文prompt的末尾准备下一次推理。

这里以GPT自回归的工作方式为例：

每步推理都将前文整句输入模型是一种效率低下的方式，原因是存在相同结果的重复推理。令前一次待推理的文本长度为S，下一次为S+1，由于网络中的各项参数已经固定，因此两次推理对于前S个token的计算结果是完全相同的， 包括Embedding映射，每一层、每一个注意力头下的KQV映射，注意力权重，以及后续的FFN层都在重复计算。

根据shift-right的性质，下一个token是由当前最后一个token的网络输出所决定的，那能不能仅输入最后一个token来进行推理？答案是否定的，虽然在结果层仅由最后一个token来决定，但是中间的注意力过程它依赖于前文所提供的Key、Value向量来携带前文信息，因此也不能抛弃前文不管。

next token计算依赖过程图

结合以上结论，S+1位置token的推理依赖于两个要素，首先是当前第S个token在网络中完整forward一遍，其次是除最后一个token以外，之前所有的S-1位置的token在每一层、每个注意力头下的Key，Value信息。又已知S-1的每个token的Key，Value信息都是在重复计算，每次计算的结果是相同的，在之前的推理中都计算过但在结果层丢弃了，因此完全可以将Key，Value信息在内存中存储起来，使得它们可以在之后的每步推理中进行复用，这种策略就是KV-Cache。这种方式避免了重复计算，大幅减少了参数的计算量，提高了推理效率。

下面我们来对比一下 Without KV-Cache 和 With KV-Cache 的区别：

• Without KV-Cache，：每次需要计算全Wq(X), Wk(X), Wv(X), 每次需要计算全量Attn。
  
• With KV-Cache：第一步计算完整Attn，将KV保存成KV_cache。第二步，取第一步的Next Token $x_n$ 计算 $Q=Wq(x_n),K=Wk(x_n),V=Wv(x_n)$， 将 [KV_cache, KV] 拼接，计算出QKV。KV-Cache每个循环累增，memory量=2*(N层*L长度*D维度)。
  
  使用KV-Cache就是缓存前面已经有的KV的tokens（缓存的是K=Wk*X），减少了X先前已经计算过了token再与Wk相乘。每次只将X中新的token与W计算得到。从下面的图中可以清晰看出整个cache的过程：

Attention with KV-Cache源码：不断在max_seq_len维度上append新的KV

class Attention(nn.Module):
    """Multi-head attention module."""
    def __init__(self, args: ModelArgs):
    
        self.wq = ColumnParallelLinear(args.dim, args.n_heads * self.head_dim)
        self.wk = ColumnParallelLinear(args.dim, self.n_kv_heads * self.head_dim)
        self.wv = ColumnParallelLinear(args.dim, self.n_kv_heads * self.head_dim)
        self.wo = RowParallelLinear(args.n_heads * self.head_dim, args.dim)
        
        
        # [8, 1024, 32, 128]
        self.cache_k = torch.zeros(
            (
                args.max_batch_size,    # 8
                args.max_seq_len,       # 1024, 不断地在这个维度上append keys
                self.n_local_kv_heads,  # 32
                self.head_dim,          # 128
            )
        ).cuda()
        
        #  [8, 1024, 32, 128]
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,    # 8
                args.max_seq_len,       # 1024, 不断地在这个维度上append values
                self.n_local_kv_heads,  # 32
                self.head_dim,          # 128
            )
        ).cuda()
    
    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        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
        
        # 这里在复用之前的计算, all_past
        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

3.6 Grouped Multi-Query Attention

Grouped-Query Attention (GQA) 是对 Multi-Head Attention (MHA) 和 Multi-Query Attention (MQA) 的扩展。通过提供计算效率和模型表达能力之间的灵活权衡，实现了查询头的分组。GQA将查询头分成了G个组，每个组共享一个公共的键（K）和值（V）投影。

GQA的优势：使用G个组可以减少存储每个头的键和值所需的内存开销，特别是在具有大的上下文窗口或批次大小的情况下。GQA提供了对模型质量和效率的细致控制。

以LLM Foundry 为例，分组查询注意力实现代码如下，与LLM Foundry 中实现的多头自注意力代码相对比，其区别仅在于建立Wqkv 层上：

import torch
import torch.nn as nn
from torch.nn import functional as F
from typing import Optional


class MultiQueryAttention(nn.Module):
    """Multi-Query self attention.
    Using torch or triton attention implemetation enables user to also use
    additive bias.
    """
    def __init__(
        self,
        d_model: int,
        n_heads: int,
        device: Optional[str] = None,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.Wqkv = nn.Linear( # Multi-Query Attention 创建
            d_model,
            d_model + 2 * self.head_dim, # 只创建查询的头向量，所以只有1 个d_model
            device=device, # 而键和值则共享各自的一个head_dim 的向量
        )
        self.out_proj = nn.Linear(
            self.d_model,
            self.d_model,
            device=device
        )
        self.out_proj._is_residual = True # type: ignore
    def forward(
        self,
        x,
    ):
        qkv = self.Wqkv(x) # (1, 512, 960)
        query, key, value = qkv.split( # query -> (1, 512, 768)
            [self.d_model, self.head_dim, self.head_dim], # key -> (1, 512, 96)
            dim=2 # value -> (1, 512, 96)
        )
        context, attn_weights, past_key_value = F.scaled_dot_product_attention(query, key, value,
            self.n_heads, multiquery=True)
        return self.out_proj(context), attn_weights, past_key_value

四、源码解析

[待补充]

参考资料

• LLaMa-1/2/3 原理+源码——拆解 (KV-Cache, RoPE, RMSNorm, GQA, SwiGLU)