探索和构建 LLaMA 3 架构：深入探究组件、编码和推理技术

最新推荐文章于 2024-08-20 21:52:32 发布

拉达曼迪斯II

最新推荐文章于 2024-08-20 21:52:32 发布

阅读量1.2k

点赞数 25

分类专栏： AI创业 AIGC学习文章标签： llama 架构人工智能网络 ai

本文链接：https://blog.csdn.net/ms44/article/details/141280676

版权

AIGC学习同时被 2 个专栏收录

144 篇文章 0 订阅

订阅专栏

AI创业

139 篇文章 0 订阅

订阅专栏

Meta 正在加大在人工智能 (AI) 竞赛中的力度，推出了新的开源 AI 模型 Llama 3 以及新版 Meta AI。这款虚拟助手由 Llama 3 提供支持，现已在所有 Meta 平台上可用。

以下是您需要了解的有关 Meta 最新大型语言模型 (LLM) 和 AI 助手的所有信息。

什么是 Llama 3？

Meta 推出了 Llama 3，这是其 Llama 系列开源 AI 模型中的最新产品。Llama 3 有两种版本：一种具有 80 亿个参数，另一种具有 700 亿个参数。

Meta 声称 Llama 3 在这些参数规模上为大型语言模型树立了新标准。他们改进了训练前和训练后流程，从而降低了错误拒绝率、提高了对齐效果，并提高了模型的响应多样性。值得注意的是，Llama 3 在推理、代码生成和指令跟踪方面的能力得到了增强。

技术规格

LLaMA 建筑事务所：

与前代模型之间的关键区别在于，预训练语料库的大小增加了 650%（LLaMA — 2 在 2T 标记上进行训练，而 LLaMA — 3 在 15T 标记上进行训练），在 8B 和 70B 模型上都将模型的上下文长度从 4K 增加了一倍至 8K，并且与上一代相比，8B 和 70B 变体都采用了分组查询注意（GQA）仅在更大的模型 34B 和 70B 中使用。我认为影响最大的部分是新的安全方法，包括针对安全性和帮助性的两种奖励模型。

LLaMA 3 吸收了上一代车型的架构

火焰 3：

模型大小、架构、优化超参数

Meta llama GitHub：llama3/MODEL_CARD.md at main·meta-llama/llama3 (github.com)

火焰2：

模型大小、架构、优化超参数

火焰 1：

LLaMA 建筑

LLaMA 3 架构主要吸收了与 LLaMA 2 相同的架构，8B 和 70B 模型中均使用了 GQA（分组查询注意），RoPE（旋转位置嵌入）用于 Q、K，因为 V 仅在应用 SoftMax 函数之前相乘，RMS（均方根误差）用于在自我注意之前应用的规范化，前馈块、KV 缓存也与 LLMA 中使用的相同

注意：该模型架构仅专注于推理模型，而不是用于训练，因此不包括具有交叉注意的解码器块，并且 KV 缓存不会用于模型的训练阶段。

现在让我们更深入地了解每个组件。

RoPE（旋转位置编码）

在深入研究 RoPE 之前，了解绝对位置编码和相对编码之间的区别非常重要。

绝对位置编码是添加到标记嵌入中的固定向量，以表示其在句子中的绝对位置。因此，它一次处理一个标记。您可以将其视为地图上的一对（纬度，经度）：地球上的每个点都会有一对唯一的标记。

另一方面，相对位置编码每次处理两个标记，并且在我们计算注意力时会涉及它：由于注意力机制捕获两个单词彼此相关的“强度”，因此相对位置编码会告诉注意力机制其中涉及的两个单词之间的距离。因此，给定两个标记，我们创建一个表示它们距离的向量。

旋转位置编码可以被认为是绝对位置嵌入和相对位置嵌入之间的中间点，因为每个标记都有一个固定或绝对的嵌入值，并且与其极坐标形式的内部点积相乘，该极坐标形式相对于二维平面上向量的旋转。

注意力机制中用到的点积是内积的一种，可以看作是点积的泛化。
我们能否找到注意力机制中使用的两个向量q（查询）和k （键）的内积，该内积仅取决于这两个向量以及它们所代表的标记的相对距离？

我们可以定义如下所示的函数g，它仅取决于两个嵌入向量q和k及其相对距离。

利用欧拉公式，我们可以将其写成矩阵形式。

二维空间中的旋转矩阵，因此称为旋转位置嵌入

来自Wolfram MathWorld：Rotation Matrix -- from Wolfram MathWorld

旋转位置嵌入仅适用于查询和键，而不适用于值。
旋转位置嵌入应用于向量问和钾已经乘以在矩阵在注意力机制中，而在 vanilla Transformer 中它们之前就被应用过了。

def  precomputed_theta_pos_frequencies ( head_dim: int , seq_len: int , device: str , theta: float = 10000.0 ): 
    # 根据论文中写到，embedding的维度必须是偶数
    assert head_dim % 2 == 0 , "head_dim 必须是偶数" 
    # 构建 theta 参数
    # 根据公式 theta_i = 10000 ^ (-2(i-1)/dim) for i = [1,2,3,..dim/2] 
    # 形状：(head_dim / 2)
     theta_numerator = torch.arange( 0 , head_dim, 2 ). float () 
    # 形状：(head_dim / 2)
     theta = 1.0 / (theta ** (theta_numerator / head_dim)).to(device) 
    # 构造位置（“m”参数）
    # 形状：(seq_len)
     m = torch.arange(seq_len, device=device) 
    # 使用外积将每个 theta 乘以每个位置
    # 形状：(seq_len) outer_product * (head_dim / 2) -> (seq_len, head_dim / 2)
     freq = torch.outer(m, theta). float () 
    # 我们可以用极坐标形式计算复数 c = R * exp(i * m * theta)，其中 R = 1，如下所示
    # 形状：(seq_len, head_dim/2) -> (seq-len, head_dim/2)
     freq_complex = torch.polar(torch.ones_like(freq), freq) 
    return freq_complex 

def  apply_rotary_embeddings ( x: torch.Tensor, freq_complex: torch.Tensor, device: str ): 
    # 我们将每个后续的标记对转换为一对复数
    # 形状：(B, seq_len, head_dim) -> (B, seq_len, h, head_dim / 2)
     x_complex = torch.view_as_complex(x. float ().reshape(*x.shape[:- 1 ], - 1 , 2 )) 
    # 形状： (seq_len, head_dim / 2) -> (1, seq_len, 1, head_dim / 2)
     freq_complex = freq_complex.unsqueeze( 0 ).unsqueeze( 2 ) 
    # 形状：(B, seq_len, h, head_dim / 2) * (1, seq_len, 1, head_dim / 2) = (B, seq_len, h, head_dim / 2)
     x_rotate = x_complex * freq_complex 
    # (B, seq_len, h, head_dim / 2) -> (B, seq_len, h, head_dim/2 ,2)
     x_out = torch.view_as_real(x_rotate) 
    # (B, seq_len, h, head_dim/2, 2) -> (B, seq_len, h * head_dim / 2 * 2)
     x_out = x_out.reshape(*x.shape)
    返回x_out.type_as(x).to(device)

键值缓存

在推理的每一步，我们只对模型输出的最后一个标记感兴趣，因为我们已经有了之前的标记。然而，模型需要访问所有之前的标记来决定输出哪个标记，因为它们构成了它的上下文（或“提示”）。
这是一种让模型在推理过程中对已经见过的 token 进行更少计算的方法。解决方案就是KV 缓存！

def  repeat_kv（x：torch.Tensor，n_rep：int）-> torch.Tensor：
    batch_size，seq_len，n_kv_heads，head_dim = x.shape 
    if n_rep == 1：
        返回x 
    else：
        返回（
            #（B，seq_len，n_kv_heads，1，head_dim）
             x [：，：，：，无，：] 
            .expand（batch_size，seq_len，n_kv_heads，n_rep，head_dim）
            .reshape（batch_size，seq_len，n_kv_heads * n_rep，head_dim）
        ）

分组多查询注意力机制

分组查询注意力 (GQA) 是多查询和多头注意力的插值。它实现了与多头注意力相似的质量，同时保持与多查询注意力相当的速度。自回归解码的标准做法是缓存序列中前一个标记的键和值以加快注意力计算。然而，随着上下文窗口或批处理大小的增加，与多头注意力 (MHA) 模型中的键值缓存 (kv 缓存) 大小相关的内存成本显著增加。多查询注意力 (MQA) 是一种仅使用单个键值头进行多个查询的机制，它可以节省内存并大大加快解码器推理速度。Llama 结合了 (GQA) 来解决 Transformer 模型自回归解码期间的内存带宽挑战。主要问题源于 GPU 执行所有计算的速度比它们将它们移动到内存中的速度更快。需要在每个阶段加载解码器权重和注意力键，这会消耗过多的内存。

class  SelfAttention (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 
        # 表示查询的头部数量
        self.n_heads_q = args.n_heads 
        # 表示键和值的头部应重复多少次才能与查询的头部匹配
        self.n_rep = self.n_heads_q // self.n_kv_heads 
        # 表示每个头部的维度
        self.head_dim = args.dim // args.n_heads 

        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias= False ) 
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim，偏差= False）
        self.wv = nn.Linear（args.dim，self.n_kv_heads * self.head_dim，偏差= False）
        self.wo = nn.Linear（args.n_heads * self.head_dim，args.dim，偏差= False）

        self.cache_k = torch.zeros（（args.max_batch_size，args.max_seq_len，self.n_kv_heads，self.head_dim））
        self.cache_v = torch.zeros（（args.max_batch_size，args.max_seq_len，self.n_kv_heads，self.head_dim））

    def  forward（self，x：torch.Tensor，start_pos：int，freq_complex：torch.Tensor）：
        batch_size， seq_len, _ = x.shape #(B, 1, dim) 
        # 将 wq、wk、wv 矩阵应用于查询、键和值
        # (B, 1, dim) -> (B, 1, H_q * head_dim)
         xq = self.wq(x) 
        # (B, 1, dim) -> (B, 1, H_kv * head_dim)
         xk = self.wk(x) 
        xv = self.wv(x) 

        # (B, 1, H_q * head_dim) -> (B, 1, H_q, head_dim)
         xq = xq.view(batch_size, seq_len, self.n_heads_q, self.head_dim) 
        xk = xk.view(batch_size, seq_len, self.n_kv_heads, self.head_dim) 
        # (B, 1, H_kv * head_dim) -> (B, 1, H_kv, head_dim)
         xv = xv.view(batch_size, seq_len, self.n_kv_heads, self.head_dim) 

        # 将旋转嵌入应用于键和值
        # 不改变张量的形状
        # (B, 1, H_kv, head_dim) -> (B, 1, H_kv, head_dim)
         xq = apply_rotary_embeddings(xq, freq_complex, device=x.device) 
        xk = apply_rotary_embeddings(xk, freq_complex, device=x.device) 

        # 用此 token 替换缓存中的实体
        self.cache_k[:batch_size, start_pos:start_pos + seq_len] = xk 
        self.cache_v[:batch_size, start_pos:start_pos + seq_len] = xv

        # 检索迄今为止所有缓存的键和值
        # (B, seq_len_kv, H_kv, head_dim)
         keys = self.cache_k[:batch_size, 0 :start_pos + seq_len] 
        values = self.cache_v[:batch_size, 0 :start_pos+seq_len] 

        # 重复 K 和 V 的头以达到查询的头数
        keys = repeat_kv(keys, self.n_rep) 
        values = repeat_kv(values, self.n_rep) 

        # (B, 1, h_q, head_dim) --> (b, h_q, 1, head_dim)
         xq = xq.transpose( 1 , 2 ) 
        keys = keys.transpose( 1 , 2 ) 
        values = values.transpose( 1 , 2 ) 

        # (B, h_q, 1, head_dim) @ (B, h_kv，seq_len-kv，head_dim）->（B，h_q，1，seq_len-kv）分数
        =torch.matmul（xq，keys.transpose（2，3 ）） / math.sqrt（self.head_dim）
        分数=F.softmax（scores.float （），dim = -1）。type_as（xq）#（B，h_q，1，seq_len）@（B，h_q，seq_len-kv，head_dim）-->（b，hq，q，head_dim）        输出=torch.matmul（scores，values）#（B，h_q，1，head_dim）->（B，1，h_q，head_dim）->（）        输出=（output.transpose（1，2 ）。contiguous（）。view（batch_size， seq_len, - 1））返回self.wo（输出）#（B，1，dim）->（B，1，dim）

RMS（均方根归一化）

均方根归一化 (RMSNorm) 是一种相对较新的归一化技术，由 Biao Zhang 和 Rico Sennrich 于 2019 年推出。与 BN 和 LN 不同，RMSNorm 根据激活本身的均方根对激活进行归一化，而不是使用小批量或层统计数据。这种方法可确保无论小批量大小或特征数量如何，激活都始终按比例缩放。此外，RMSNorm 引入了可学习的缩放参数，提供与批量归一化类似的适应性。

注意：与层规范化一样，我们也有一个可学习的参数gamma（左边公式中的g ），它与规范化值相乘。

class  RMSNorm (nn.Module): 
    def  __init__ ( self, dim: int , eps: float = 1e-5 ): 
        super ().__init__() 
        self.eps = eps 
        # gamma 参数
        self.weight = nn.Parameter(torch.ones(dim)) 

    def  _norm ( self, x: torch.Tensor ): 
        # (B, seq_len, dim) -> (B, seq_len, 1) 
        return x * torch.rsqrt(x.pow ( 2 ) .mean(- 1 , keepdim= True ) + self.eps) 

    def  forward ( self, x: torch.Tensor ): 
        # dim : (B, seq_len, dim) -> (B, seq_len, dim) 
        return self.weight * self._norm(x. float ()).type_as(x)

SwiGLU 激活函数

SwiGLU 是深度神经网络中使用的一种激活函数，是 GLU（门控线性单元）的一种变体。它用于通过获取输入的加权和并对其应用非线性函数来计算神经网络中神经元的输出。SwiGLU 使用涉及 Swish 函数和张量乘法的数学表达式来定义。SwiGLU 是 GLU 的一种变体，这意味着它基于与 GLU 相同的数学概念。但是，SwiGLU 具有与 GLU 不同的非线性函数。具体来说，SwiGLU 使用 Swish 函数，这是一种最近提出的激活函数，在某些应用中已被证明优于其他激活函数。

SwiGLU 具有多种优势，使其成为神经网络中有用的激活函数。首先，它基于 GLU 概念，该概念已被证明在许多应用中表现良好。其次，它使用 Swish 函数，该函数已被证明在某些情况下优于其他激活函数，特别是与残差连接结合使用时。第三，由于它使用元素乘法，因此可以实现高效计算。

作者通过在 Transformer 架构的前馈层中使用不同的激活函数来比较 Transformer 模型的性能。

def  forward ( self, x: torch.Tensor ): 
        swish = F.silu(self.w1(x))   # 应用第一个变换
        x_V = self.w3(x) 
        x = swish * x_V         # 对原始维度应用收缩
        x = self.w2(x)   # 应用可选的附加变换
        return x

前馈模块

在 Transformer 架构中，前馈层起着至关重要的作用，通常位于注意层和规范化之后。前馈层由三个线性变换组成。

class  FeedForward (nn.Module): 
    def  __init__ ( self, args: ModelArgs ): 
        super ().__init__() 
        # 假设 'hidden_dim' 是根据您的规范计算的
        hidden_dim = 4 * args.dim 
        hidden_dim = int ( 2 * hidden_dim / 3 )   # 应用您的特定转换
        if args.ffn_dim_multiplier不是 None :             hidden_dim = int (args.ffn_dim_multiplier * hidden_dim) #hidden_ dim  = int(2 * hidden_dim / 3) # 应用您的特定转换        hidden_dim = args.multiple_of * ((hidden_dim + args.multiple_of - 1 ) // args.multiple_of)         self.w1 = nn.Linear(args.dim, hidden_dim, bias= False )         self.w2 = nn.Linear(hidden_dim, args.dim, bias= False )   # 您的原始设置中似乎缺少这一层        self.w3 = nn.Linear(args.dim, hidden_dim, bias= False )   # 已更正以匹配检查点def forward ( self, x: torch.Tensor ):         swish = F.silu(self.w1(x))   # 应用第一个变换        x_V = self.w3(x)         x = swish * x_V         # 对原始维度应用收缩        x = self.w2(x)   # 应用可选的附加变换return x

在前向传递过程中，输入张量x会经历多层线性变换。在第一次变换后应用的SwiGLU激活函数可增强模型的表达能力。最终的变换将张量映射回其原始维度。SwiGLU 激活和多个前馈层的独特组合可增强模型的性能。

编码器模块

由于我们只关注模型的推理，所以我们只研究编码器块

class  EncoderBlock (nn.Module): 
    def  __init__ ( self, args: ModelArgs ): 
        super ().__init__() 
        self.n_heads = args.n_heads 
        self.dim = args.dim 
        self.head_dim = args.dim // args.n_heads 

        self.attention = SelfAttention(args) 
        self.feed_forward = FeedForward(args) 

        # 在自我注意之前进行标准化
        self.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_complex: torch.Tensor ): 
        # (B, seq_len, dim) + (B, seq_len，dim）->（B，seq_len，dim）
         h = x + self.attention.forward（self.attention_norm（x），start_pos，freqs_complex）
        out = h + self.feed_forward.forward（self.ffn_norm（h））
        返回out

终极变压器模型

现在让我们把所有这些单独的组件堆叠到 Transformer Block 中

class  ModelArgs : 
    dim: int = 4096
     n_layers: int = 32
     n_heads: int = 32  # 查询的 head 数量
    n_kv_heads: Optional [ int ] = None  # 键和值的 head 数量。如果为 None，则默认为 n_heads
     vocab_size: int = - 1  # 加载 tokenizer 时会设置
    multiple_of: int = 256 
     ffn_dim_multiplier: Optional [ float ] = None  # 如果为 None，则默认为 4.0
     norm_eps: float = 1e-6  # 仅修改了 eps 值 int llama 3 
    
    # KV 缓存所需
    max_batch_size: int = 32
     max_seq_len: int = 2048

     device: str = None

class  Transformer (nn.Module): 
    
    def  __init__ ( self, args: ModelArgs ) -> None : 
        super ().__init__() 

        assert args.vocab_size != - 1 , "必须设置词汇大小"

         self.args = args 
        self.vocab_size = args.vocab_size 
        self.n_layers = args.n_layers 
        self.tok_embeddings = nn.Embedding(self.vocab_size, args.dim) 

        self.layers = nn.ModuleList() 
        for _ in  range (args.n_layers): 
            self.layers.append(EncoderBlock(args)) 

        self.norm = RMSNorm(args.dim, eps=args.norm_eps) 
        self.output = nn.Linear(args.dim, self.vocab_size, bias= False ) 

        # 预先计算频率旋转位置编码
        self.freqs_complex = precomputed_theta_pos_frequencies(self.args.dim // self.args.n_heads, self.args.max_seq_len * 2 , device=self.args.device) 

    def  forward ( self, tokens: torch.Tensor, start_pos: int ): 
        # (B, seq_len)
         batch_size, seq_len = tokens.shape 
        assert seq_len == 1 , "一次只能处理一个 token" 
  
        # (B, seq_len) -> (B, seq_len, dim)
         h = self.tok_embeddings(tokens) 

        # 检索与位置 [start_pos, start_pos + seq_len] 相对应的对 (m, theta)
         freqs_complex = self.freqs_complex[start_pos:start_pos + seq_len] 

        #连续应用self.layers中的
        所有编码器层：            h = layer(h, start_pos, freqs_complex)         h = self.norm(h)         output = self.output(h)。float ()返回输出

推理

为了推理模型，我们需要从 Meta 的 Llama 3 repo 下载模型的权重。

拉达曼迪斯II

关注

25
点赞
踩
22

收藏

觉得还不错? 一键收藏
打赏
0
评论
探索和构建 LLaMA 3 架构：深入探究组件、编码和推理技术

Meta 推出了 Llama 3，这是其 Llama 系列开源 AI 模型中的最新产品。Llama 3 有两种版本：一种具有 80 亿个参数，另一种具有 700 亿个参数。Meta 声称 Llama 3 在这些参数规模上为大型语言模型树立了新标准。他们改进了训练前和训练后流程，从而降低了错误拒绝率、提高了对齐效果，并提高了模型的响应多样性。值得注意的是，Llama 3 在推理、代码生成和指令跟踪方面的能力得到了增强。© 由 Meta 提供/wp:image。
复制链接

扫一扫