LLaMa系列模型详解(原理介绍、代码解读):LLaMA 2

LLaMA 2

大型语言模型(LLMs)作为高度能力的人工智能助手,在需要跨多个领域专家知识的复杂推理任务中表现出巨大潜力,包括编程和创意写作等专业领域。它们通过直观的聊天界面与人类互动,这导致了快速和广泛的公众采用。考虑到训练方法的看似简单性,LLMs 的能力令人瞩目。自回归变压器首先在大量自监督数据上进行预训练,然后通过强化学习与人类反馈(RLHF)等技术与人类偏好对齐。尽管训练方法简单,但高计算需求限制了 LLMs 的开发,仅由少数参与者进行。

image.png

虽然有一些预训练的 LLMs 公开发布(如 BLOOM、LLaMa-1 和 Falcon),它们在性能上可以匹敌封闭的预训练竞争对手(如 GPT-3 和 Chinchilla),但这些模型都不适合作为封闭“产品”LLMs(如 ChatGPT、BARD 和 Claude)的替代品。这些封闭的产品 LLMs 被大量微调以对齐人类偏好,这大大提高了它们的可用性和安全性。这个步骤可能需要大量的计算成本和人类注释,并且往往不透明或不易复制,限制了社区在 AI 对齐研究方面的进展。

image.png

为了推动这一领域的发展,Llama 2 系列模型被开发并发布,包括预训练和微调的 LLMs:Llama 2 和 Llama 2-Chat,参数规模高达 70B。在一系列有用性和安全性基准上,Llama 2-Chat 模型通常表现优于现有的开源模型。根据人类评估结果,这些模型似乎也与一些封闭源模型相当。Llama 2 系列模型通过安全特定的数据注释和调优,以及红队测试和迭代评估,增加了其安全性。此外,这些模型的微调方法和改进 LLM 安全性的途径也被详细描述。希望这种开放性能够使社区复制微调的 LLMs 并继续提高这些模型的安全性,为更负责任的 LLM 开发铺平道路。在开发 Llama 2 和 Llama 2-Chat 过程中,还发现了一些新奇的现象,例如工具使用的出现和知识的时间组织。

以下模型已向公众发布,用于研究和商业使用:

  1. Llama 2,Llama 1 的更新版本,训练于新的公开数据混合体上。预训练语料库的大小增加了 40%,模型的上下文长度翻倍,并采用了分组查询注意力。发布了 7B、13B 和 70B 参数的 Llama 2 变体。
  2. Llama 2-Chat,Llama 2 的微调版本,针对对话用例进行了优化。发布了 7B、13B 和 70B 参数的这个模型变体。

image.png

模型架构

采用 Llama 1 中的大部分预训练设置和模型架构。使用标准 Transformer 架构,使用 RMSNorm 应用预归一化,使用 SwiGLU 激活函数和旋转位置嵌入。与 Llama 1的主要架构差异包括增加的上下文长度和分组查询注意力 (GQA)。

image.png

本文主要介绍LLaMA 2和LLaMA 1的区别部分,如果想具体了解LLaMA 1的模型架构和代码解读请点击此处

分组查询注意力 (GQA)

增加上下文长度比较好理解,简单的在训练前规定了最大上下文长度为4096,本文主要介绍LLaMA2中改进的注意力机制。

在理解什么是GQA之前,我们还需要知道两个概念:MHA和MQA,下图展示了MHA,MQA,GQA的区别:
image.png

MHA

多头注意力机制MHA(Multi-Head Attention),将输入数据分成多个头(heads),每个头独立地执行注意力计算。这些头通常具有不同的权重矩阵,因此可以关注输入序列中的不同部分和特征。QKV 三部分有相同数量的头,且一一对应。每次做 Attention,head1 的 QKV 就做好自己运算就可以,输出时各个头加起来就行。

MQA

多查询注意力机制MQA(Multi-Query Attention),MQA的原理很简单,简单来说Q仍然是多头,K,V是共享的。它将原生Transformer每一层多头注意力的Key线性映射矩阵、Value线性映射矩阵改为该层下所有头共享,也就是说K、V矩阵每层只有一个。举例来说,以ChatGLM2-6B为例,一共28层,32个注意力头,输入维度从4096经过Q、K、V矩阵映射维度为128,若采用原生多头注意力机制,则Q、K、V矩阵各有28×32个,而采用MQA的方式则整个模型包含28×32个Q矩阵,28×1个K矩阵,28×1个V矩阵。这种方法在提高推理效率的同时,也能够保持模型的性能。

GQA

MQA虽然能最大程度减少KV Cache所需的缓存空间,但是可想而知参数的减少意味着精度的下降,所以为了在精度和计算之间做一个trade-off,GQA (Group Query Attention)应运而生,即Q依然是多头,但是分组共享K,V,既减少了K,V缓存所需的缓存空间,也暴露了大部分参数不至于精度损失严重。

KV Cache

大模型推理性能优化的一个常用技术是KV Cache,那么什么是KV Cache呢?

在自回归生成任务中,模型需要逐个生成序列中的tokens,每次生成一个新token时,都会更新输入序列并重新计算自注意力。然而,已生成的部分(历史tokens)对应的Key和Value向量在生成后续token时往往保持不变或变化较小。KV Cache正是利用了这一性质,通过将这些历史tokens对应的Key和Value向量存储起来(缓存),在后续计算中直接复用,而不是每次都重新计算。

代码详解

RMSNorm(均方根归一化)

代码实现的是对输入张量 x 进行RMS归一化,将每个元素除以其均方根(RMS),并确保计算过程的数值稳定性。

class RMSNorm(torch.nn.Module):  def __init__(self, dim: int, eps: float = 1e-6):  """  初始化 RMSNorm 归一化层。  参数:  dim (int): 输入张量的维度。  eps (float, 可选): 添加到分母的小值,以确保数值稳定性。默认值为 1e-6。  属性:  eps (float): 添加到分母的小值,以确保数值稳定性。  weight (nn.Parameter): 可学习的缩放参数。  """        super().__init__()  self.eps = eps  self.weight = nn.Parameter(torch.ones(dim))  def _norm(self, x):  """  对输入张量应用 RMSNorm 归一化。  参数:  x (torch.Tensor): 输入张量。  返回:  torch.Tensor: 归一化后的张量。  """        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)  def forward(self, x):  """  通过 RMSNorm 层的前向传递。  参数:  x (torch.Tensor): 输入张量。  返回:  torch.Tensor: 应用 RMSNorm 后的输出张量。  """        output = self._norm(x.float()).type_as(x)  return output * self.weight

旋转位置嵌入

为了便于理解,给出RoPE的具体实现步骤:

  1. 频率向量的计算:
    f i = 1 θ 2 i d f_i = \frac{1}{\theta^{\frac{2i}{d}}} fi=θd2i1
    其中 θ \theta θ是一个常数(通常取 10000), i i i是向量维度的索引。

  2. 旋转角度的计算:
    angle ( t ) = t ⋅ f i \text{angle}(t) = t \cdot f_i angle(t)=tfi
    其中 t t t是位置索引。

  3. 应用旋转变换:
    对每个位置 t t t的输入向量 x t x_t xt,在复数域进行旋转变换:
    x t ′ = x t ⋅ e j ⋅ angle ( t ) x_t' = x_t \cdot e^{j \cdot \text{angle}(t)} xt=xtejangle(t)

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):  """  预计算复数指数的频率张量(cis),具有给定的维度。  此函数使用给定的维度 'dim' 和结束索引 'end' 计算一个复数指数的频率张量。'theta' 参数用于缩放频率。  返回的张量包含复数值,数据类型为 complex64。  参数:  dim (int): 频率张量的维度。  end (int): 用于预计算频率的结束索引。  theta (float, 可选): 用于频率计算的缩放因子。默认为 10000.0。  返回:  torch.Tensor: 预计算的复数指数频率张量。  """    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)  # complex64  return freqs_cis

代码相当于预先了计算了angle[t]列表,将每个位置的旋转矩阵保存下来,减少训练中的计算。

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):  """  重塑频率张量以便与另一个张量进行广播。  此函数将频率张量重塑为与目标张量 'x' 具有相同的形状,以便在进行逐元素操作时进行广播。 参数:  freqs_cis (torch.Tensor): 需要重塑的频率张量。  x (torch.Tensor): 目标张量,用于广播兼容性。  返回:  torch.Tensor: 重塑后的频率张量。  抛出:  AssertionError: 如果频率张量的形状不符合预期。  AssertionError: 如果目标张量 'x' 没有预期的维数。  """    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)

假设

  • freqs_cis 的形状为 (L, D),其中 L 是序列长度,D 是特征维度。
  • x 的形状为 (B, L, H, D),其中 B 是批量大小,L 是序列长度,H 是头数,D 是每个头的特征维度。
    结果
  • 将频率张量 freqs_cis 重塑为 [1, L, 1, D],这个形状可以和x进行广播。
def apply_rotary_emb(  xq: torch.Tensor,  xk: torch.Tensor,  freqs_cis: torch.Tensor,  
) -> Tuple[torch.Tensor, torch.Tensor]:  """  使用给定的频率张量对输入张量应用旋转嵌入。  此函数使用提供的频率张量 'freqs_cis' 对给定的查询 'xq' 和键 'xk' 张量应用旋转嵌入。  输入张量被重塑为复数,并重塑频率张量以进行广播兼容性。返回的张量包含旋转嵌入,并以实数形式返回。  参数:  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))  xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  freqs_cis = reshape_for_broadcast(freqs_cis, xq_)  xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)  xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)  return xq_out.type_as(xq), xk_out.type_as(xk)

函数实现了在Attention中使用旋转位置嵌入,与一般位置嵌入不同,旋转位置嵌入在计算完QK后进行,直接在QK上增加旋转位置信息。

应用KV-Cache和GQA的注意力机制

class Attention(nn.Module):  """多头注意力模块。"""  def __init__(self, args: ModelArgs):  """  初始化 Attention 模块。  参数:  args (ModelArgs): 模型配置参数。  属性:  n_kv_heads (int): 键和值的头数。  n_local_heads (int): 本地查询头数。  n_local_kv_heads (int): 本地键和值头数。  n_rep (int): 本地头的重复次数。  head_dim (int): 每个注意力头的维度大小。  wq (ColumnParallelLinear): 查询的线性变换。  wk (ColumnParallelLinear): 键的线性变换。  wv (ColumnParallelLinear): 值的线性变换。  wo (RowParallelLinear): 输出的线性变换。  cache_k (torch.Tensor): 注意力的缓存键。  cache_v (torch.Tensor): 注意力的缓存值。  """        super().__init__()  self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads  model_parallel_size = fs_init.get_model_parallel_world_size()  self.n_local_heads = args.n_heads // model_parallel_size  self.n_local_kv_heads = self.n_kv_heads // model_parallel_size  self.n_rep = self.n_local_heads // self.n_local_kv_heads  self.head_dim = args.dim // args.n_heads  self.wq = ColumnParallelLinear(  args.dim,  args.n_heads * self.head_dim,  bias=False,  gather_output=False,  init_method=lambda x: x,  )  self.wk = ColumnParallelLinear(  args.dim,  self.n_kv_heads * self.head_dim,  bias=False,  gather_output=False,  init_method=lambda x: x,  )  self.wv = ColumnParallelLinear(  args.dim,  self.n_kv_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_kv_heads,  self.head_dim,  )  ).cuda()  self.cache_v = torch.zeros(  (  args.max_batch_size,  args.max_seq_len,  self.n_local_kv_heads,  self.head_dim,  )  ).cuda()  def forward(  self,  x: torch.Tensor,  start_pos: int,  freqs_cis: torch.Tensor,  mask: Optional[torch.Tensor],  ):  """  Attention 模块的前向传递。  参数:  x (torch.Tensor): 输入张量。  start_pos (int): 缓存的起始位置。  freqs_cis (torch.Tensor): 预计算的频率张量。  mask (torch.Tensor, 可选): 注意力掩码张量。  返回:  torch.Tensor: 注意力后的输出张量。  """        bsz, seqlen, _ = x.shape  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  keys = self.cache_k[:bsz, : start_pos + seqlen]  values = self.cache_v[:bsz, : start_pos + seqlen]  # 如果 n_kv_heads < n_heads,则重复 k/v heads        keys = repeat_kv(keys, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)  values = repeat_kv(values, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)  xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)  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, seqlen, cache_len + seqlen)  scores = F.softmax(scores.float(), dim=-1).type_as(xq)  output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)  output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)  return self.wo(output)

KV缓存机制

代码定义了KV缓存机制:

self.cache_k = torch.zeros( ( args.max_batch_size, args.max_seq_len, self.n_local_kv_heads, self.head_dim, ) ).cuda() self.cache_v = torch.zeros( ( args.max_batch_size, args.max_seq_len, self.n_local_kv_heads, self.head_dim, ) ).cuda()

具体实现是,每次计算完KV后,将本次计算结果加入cache_k和cache_v后:

self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk  
self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

GQA的实现

初始化头的分割

self.n_local_heads = args.n_heads // model_parallel_size
self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
self.n_rep = self.n_local_heads // self.n_local_kv_heads

n_local_heads是Q的头数,n_local_kv_heads是KV的头数,n_rep是为每个KV头的重复次数。

变换和重塑

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)

重塑为对应的形状

键和值的重复使用

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:  """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""  bs, slen, n_kv_heads, head_dim = x.shape  if n_rep == 1:  return x  return (  x[:, :, :, None, :]  .expand(bs, slen, n_kv_heads, n_rep, head_dim)  .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  )

这里,键和值被重复以匹配查询的数量,确保每组中的查询都有相应的键和值可用。

keys = repeat_kv(keys, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
values = repeat_kv(values, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)

注意力分数计算与应用

scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)

整体实现

上面解释了LLaMA系列对Transformer的改进,下面给出全部的LLaMA模型实现:

注:LLaMA系列统一对输入进行归一化,而不是对输出进行归一化。

# 版权所有 (c) Meta Platforms, Inc. 及其子公司。  
# 本软件可以根据 Llama 2 社区许可协议的条款使用和分发。  import math  
from dataclasses import dataclass  
from typing import Optional, Tuple  import fairscale.nn.model_parallel.initialize as fs_init  
import torch  
import torch.nn.functional as F  
from fairscale.nn.model_parallel.layers import (  ColumnParallelLinear,  ParallelEmbedding,  RowParallelLinear,  
)  
from torch import nn  @dataclass  
class ModelArgs:  dim: int = 4096  n_layers: int = 32  n_heads: int = 32  n_kv_heads: Optional[int] = None  vocab_size: int = -1  # 稍后由 tokenizer 定义  multiple_of: int = 256  # 使 SwiGLU 隐藏层大小成为大的2的幂的倍数  ffn_dim_multiplier: Optional[float] = None  norm_eps: float = 1e-5  max_batch_size: int = 32  max_seq_len: int = 2048  class RMSNorm(torch.nn.Module):  def __init__(self, dim: int, eps: float = 1e-6):  """  初始化 RMSNorm 归一化层。  参数:  dim (int): 输入张量的维度。  eps (float, 可选): 添加到分母的小值,以确保数值稳定性。默认值为 1e-6。  属性:  eps (float): 添加到分母的小值,以确保数值稳定性。  weight (nn.Parameter): 可学习的缩放参数。  """        super().__init__()  self.eps = eps  self.weight = nn.Parameter(torch.ones(dim))  def _norm(self, x):  """  对输入张量应用 RMSNorm 归一化。  参数:  x (torch.Tensor): 输入张量。  返回:  torch.Tensor: 归一化后的张量。  """        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)  def forward(self, x):  """  通过 RMSNorm 层的前向传递。  参数:  x (torch.Tensor): 输入张量。  返回:  torch.Tensor: 应用 RMSNorm 后的输出张量。  """        output = self._norm(x.float()).type_as(x)  return output * self.weight  def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):  """  预计算复数指数的频率张量(cis),具有给定的维度。  此函数使用给定的维度 'dim' 和结束索引 'end' 计算一个复数指数的频率张量。'theta' 参数用于缩放频率。  返回的张量包含复数值,数据类型为 complex64。  参数:  dim (int): 频率张量的维度。  end (int): 用于预计算频率的结束索引。  theta (float, 可选): 用于频率计算的缩放因子。默认为 10000.0。  返回:  torch.Tensor: 预计算的复数指数频率张量。  """    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)  # complex64  return freqs_cis  def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):  """  重塑频率张量以便与另一个张量进行广播。  此函数将频率张量重塑为与目标张量 'x' 具有相同的形状,以便在进行逐元素操作时进行广播。  参数:  freqs_cis (torch.Tensor): 需要重塑的频率张量。  x (torch.Tensor): 目标张量,用于广播兼容性。  返回:  torch.Tensor: 重塑后的频率张量。  抛出:  AssertionError: 如果频率张量的形状不符合预期。  AssertionError: 如果目标张量 'x' 没有预期的维数。  """    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]:  """  使用给定的频率张量对输入张量应用旋转嵌入。  此函数使用提供的频率张量 'freqs_cis' 对给定的查询 'xq' 和键 'xk' 张量应用旋转嵌入。  输入张量被重塑为复数,并重塑频率张量以进行广播兼容性。返回的张量包含旋转嵌入,并以实数形式返回。  参数:  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))  xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  freqs_cis = reshape_for_broadcast(freqs_cis, xq_)  xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)  xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)  return xq_out.type_as(xq), xk_out.type_as(xk)  def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:  """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""  bs, slen, n_kv_heads, head_dim = x.shape  if n_rep == 1:  return x  return (  x[:, :, :, None, :]  .expand(bs, slen, n_kv_heads, n_rep, head_dim)  .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  )  class Attention(nn.Module):  """多头注意力模块。"""  def __init__(self, args: ModelArgs):  """  初始化 Attention 模块。  参数:  args (ModelArgs): 模型配置参数。  属性:  n_kv_heads (int): 键和值的头数。  n_local_heads (int): 本地查询头数。  n_local_kv_heads (int): 本地键和值头数。  n_rep (int): 本地头的重复次数。  head_dim (int): 每个注意力头的维度大小。  wq (ColumnParallelLinear): 查询的线性变换。  wk (ColumnParallelLinear): 键的线性变换。  wv (ColumnParallelLinear): 值的线性变换。  wo (RowParallelLinear): 输出的线性变换。  cache_k (torch.Tensor): 注意力的缓存键。  cache_v (torch.Tensor): 注意力的缓存值。  """        super().__init__()  self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads  model_parallel_size = fs_init.get_model_parallel_world_size()  self.n_local_heads = args.n_heads // model_parallel_size  self.n_local_kv_heads = self.n_kv_heads // model_parallel_size  self.n_rep = self.n_local_heads // self.n_local_kv_heads  self.head_dim = args.dim // args.n_heads  self.wq = ColumnParallelLinear(  args.dim,  args.n_heads * self.head_dim,  bias=False,  gather_output=False,  init_method=lambda x: x,  )  self.wk = ColumnParallelLinear(  args.dim,  self.n_kv_heads * self.head_dim,  bias=False,  gather_output=False,  init_method=lambda x: x,  )  self.wv = ColumnParallelLinear(  args.dim,  self.n_kv_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_kv_heads,  self.head_dim,  )  ).cuda()  self.cache_v = torch.zeros(  (  args.max_batch_size,  args.max_seq_len,  self.n_local_kv_heads,  self.head_dim,  )  ).cuda()  def forward(  self,  x: torch.Tensor,  start_pos: int,  freqs_cis: torch.Tensor,  mask: Optional[torch.Tensor],  ):  """  Attention 模块的前向传递。  参数:  x (torch.Tensor): 输入张量。  start_pos (int): 缓存的起始位置。  freqs_cis (torch.Tensor): 预计算的频率张量。  mask (torch.Tensor, 可选): 注意力掩码张量。  返回:  torch.Tensor: 注意力后的输出张量。  """        bsz, seqlen, _ = x.shape  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  keys = self.cache_k[:bsz, : start_pos + seqlen]  values = self.cache_v[:bsz, : start_pos + seqlen]  # 如果 n_kv_heads < n_heads,则重复 k/v heads        keys = repeat_kv(keys, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)  values = repeat_kv(values, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)  xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)  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, seqlen, cache_len + seqlen)  scores = F.softmax(scores.float(), dim=-1).type_as(xq)  output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)  output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)  return self.wo(output)  class FeedForward(nn.Module):  def __init__(  self,  dim: int,  hidden_dim: int,  multiple_of: int,  ffn_dim_multiplier: Optional[float],  ):  """  初始化 FeedForward 模块。  参数:  dim (int): 输入维度。  hidden_dim (int): 前馈层的隐藏维度。  multiple_of (int): 确保隐藏维度是该值的倍数。  ffn_dim_multiplier (float, 可选): 隐藏维度的自定义乘数。默认为 None。  属性:  w1 (ColumnParallelLinear): 第一层的线性变换。  w2 (RowParallelLinear): 第二层的线性变换。  w3 (ColumnParallelLinear): 第三层的线性变换。  """        super().__init__()  hidden_dim = int(2 * hidden_dim / 3)  # 自定义维度因子乘数  if ffn_dim_multiplier is not None:  hidden_dim = int(ffn_dim_multiplier * hidden_dim)  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))  class TransformerBlock(nn.Module):  def __init__(self, layer_id: int, args: ModelArgs):  """  初始化一个 TransformerBlock。  参数:  layer_id (int): 层的标识符。  args (ModelArgs): 模型配置参数。  属性:  n_heads (int): 注意力头数。  dim (int): 模型的维度大小。  head_dim (int): 每个注意力头的维度大小。  attention (Attention): 注意力模块。  feed_forward (FeedForward): 前馈模块。  layer_id (int): 层的标识符。  attention_norm (RMSNorm): 注意力输出的层归一化。  ffn_norm (RMSNorm): 前馈输出的层归一化。  """        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,  ffn_dim_multiplier=args.ffn_dim_multiplier,  )  self.layer_id = layer_id  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_cis: torch.Tensor,  mask: Optional[torch.Tensor],  ):  """  通过 TransformerBlock 的前向传递。  参数:  x (torch.Tensor): 输入张量。  start_pos (int): 注意力缓存的起始位置。  freqs_cis (torch.Tensor): 预计算的余弦和正弦频率。  mask (torch.Tensor, 可选): 注意力的掩码张量。默认值为 None。  返回:  torch.Tensor: 应用注意力和前馈层后的输出张量。  """        h = x + self.attention.forward(  self.attention_norm(x), start_pos, freqs_cis, mask  )  out = h + self.feed_forward.forward(self.ffn_norm(h))  return out  class Transformer(nn.Module):  def __init__(self, params: ModelArgs):  """  初始化一个 Transformer 模型。  参数:  params (ModelArgs): 模型配置参数。  属性:  params (ModelArgs): 模型配置参数。  vocab_size (int): 词汇表大小。  n_layers (int): 模型的层数。  tok_embeddings (ParallelEmbedding): 词嵌入。  layers (torch.nn.ModuleList): Transformer 块的列表。  norm (RMSNorm): 模型输出的层归一化。  output (ColumnParallelLinear): 最终输出的线性层。  freqs_cis (torch.Tensor): 预计算的余弦和正弦频率。  """        super().__init__()  self.params = params  self.vocab_size = params.vocab_size  self.n_layers = params.n_layers  self.tok_embeddings = ParallelEmbedding(  params.vocab_size, params.dim, init_method=lambda x: x  )  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):  """  通过 Transformer 模型的前向传递。  参数:  tokens (torch.Tensor): 输入的标记索引。  start_pos (int): 注意力缓存的起始位置。  返回:  torch.Tensor: 应用 Transformer 模型后的输出 logits。  """        _bsz, seqlen = tokens.shape  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 = None  if 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)  for layer in self.layers:  h = layer(h, start_pos, freqs_cis, mask)  h = self.norm(h)  output = self.output(h).float()  return output

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/842536.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

军工单位如何做到安全跨网文件交换与导出的

在现代信息化战争中&#xff0c;军工单位在信息安全方面的需求尤为突出。跨网文件交换与导出作为军工单位日常运营的重要环节&#xff0c;面临着网络带宽限制、数据安全风险、合规性要求和传输稳定性等挑战。下面&#xff0c;我们将从以下几个方面探讨军工单位如何实现安全、高…

【C++】开源:RabbitMQ安装与配置使用(SimpleAmqpClient)

&#x1f60f;★,:.☆(&#xffe3;▽&#xffe3;)/$:.★ &#x1f60f; 这篇文章主要介绍。 无专精则不能成&#xff0c;无涉猎则不能通。——梁启超 欢迎来到我的博客&#xff0c;一起学习&#xff0c;共同进步。 喜欢的朋友可以关注一下&#xff0c;下次更新不迷路&#x1…

用队列实现栈 用栈实现队列 设计循环队列

用队列实现栈 思路 栈的特点&#xff1a;后进先出 队列的特点&#xff1a;先进先出 使用两个队列实现栈&#xff1a; 我们可以使用两个队列&#xff0c;一个队列为&#xff1a;空队列&#xff0c;一个队列为&#xff1a;非空队列 当我们要出队列时&#xff1a; 将 size - …

多线程JUC 第2季 BlockingQueue 阻塞队列

一 阻塞队列 1.1 阻塞队列介绍 阻塞队列&#xff08;BlockingQueue&#xff09;是一个在队列基础上又支持了两个附加操作的队列&#xff1a; put方法&#xff1a;当队列装满时&#xff0c;添加的线程则被阻塞&#xff0c;直到队列不满&#xff0c;则可用。 take方法&#x…

力扣2028. 找出缺失的观测数据

题目&#xff1a; 现有一份 n m 次投掷单个 六面 骰子的观测数据&#xff0c;骰子的每个面从 1 到 6 编号。观测数据中缺失了 n 份&#xff0c;你手上只拿到剩余 m 次投掷的数据。幸好你有之前计算过的这 n m 次投掷数据的 平均值 。 给你一个长度为 m 的整数数组 rolls &a…

【Linux 网络】网络基础(三)(数据链路层协议:以太网协议、ARP 协议)

一、以太网 两个不同局域网的主机传递数据并不是直接传递的&#xff0c;而是通过路由器 “一跳一跳” 的传递过去。 跨网络传输的本质&#xff1a;由无数个局域网&#xff08;子网&#xff09;转发的结果。 所以&#xff0c;要理解数据跨网络转发原理就要先理解一个局域网中数…

安全厂商第一站!OASA 走进绿盟科技圆满结束

近日&#xff0c;龙蜥社区安全联盟&#xff08;OASA&#xff09;走进联盟成员单位绿盟科技集团股份有限公司&#xff08;以下简称“绿盟科技”&#xff09;&#xff0c;就未来合作方向&#xff0c;双方进行了一次深入的合作交流。该会议共有 11 位人员出席&#xff0c;有来自绿…

视频监控平台AS-V1000产品介绍:账户或用户数据的导入和导出功能介绍

目录 一、功能描述 &#xff08;一&#xff09;导入功能定义 &#xff08;二&#xff09;导出功能定义 二、用户数据的导入导出的作用 三、AS-V1000新版本的导出和导入功能介绍 &#xff08;一&#xff09;功能主界面 &#xff08;二&#xff09;导出功能 1、导出操作 …

从零开始构建 Vision Transformer(ViT) 模型

Transformer 模型最早由 Vaswani 等人在 2017 年论文 Attention Is All You Need 中提出&#xff0c;并已广泛应用于自然语言处理。 2021年&#xff0c;Dosovitsky 等人在论文An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale中提出将 Transforme…

第十二届蓝桥杯物联网试题(国赛)

不得不说国赛相比较省赛而言确实&#xff0c;功能变得更加复杂&#xff0c;更加繁琐&#xff0c;特别是串口LORA通信相结合的更加频繁&#xff0c;且对收取的字符处理要求要更加复杂&#xff0c;处理判别起来会更加复杂。 对于收发数据本身来说&#xff0c;收发的数据本身是以…

每日一题——Python实现PAT甲级1029 Median(举一反三+思想解读+逐步优化)

一个认为一切根源都是“自己不够强”的INTJ 个人主页&#xff1a;用哲学编程-CSDN博客专栏&#xff1a;每日一题——举一反三Python编程学习Python内置函数 Python-3.12.0文档解读 目录 我的方法 代码功能和结构点评 时间复杂度分析 空间复杂度分析 优化建议 我要更强…

深度学习环境安装教程-anaconda-python-pytorch

首先是anaconda的安装&#xff0c;可以从下面地址下载安装包 Index of /anaconda/archive/ | 清华大学开源软件镜像站 | Tsinghua Open Source Mirror 尽量选择最新的日期的anaconda进行安装&#xff0c;我这里是windows电脑&#xff0c;因此选择了windos-x86_64.exe&#xf…

如何在anaconda的环境下安装langchain

1、安装anaconda&#xff1b; 2、在终端上&#xff0c;输入&#xff1a; conda install langchain -c conda-forge Proceed ([y]/n)? y 输入&#xff1a;Y 3、安装完成后&#xff0c;输入&#xff1a; python -c "import langchain; print(langchain.__version__)&…

Redisson集成SpringBoot

前言&#xff1a;Redisson集成SpringBoot主要有两种方式&#xff0c;一个是使用redisson-spring-boot-starter依赖&#xff08;优先推荐&#xff09;&#xff0c;毕竟springboot主打的就是约定大于配置&#xff0c;这个依赖就是为springboot准备的。 再一种方式就是引入rediss…

JVM学习-javap解析Class文件

解析字节码的作用 通过反编译生成字节码文件&#xff0c;可以深入了解Java工作机制&#xff0c;但自己分析类文件结构太麻烦&#xff0c;除了第三方的jclasslib工具外&#xff0c;官方提供了javapjavap是jdk自带的反解析工具&#xff0c;它的作用是根据class字节码文件&#x…

一个投稿好方法让你的文章早日发表

作为一名单位信息宣传员,我初入此行时,满腔热情,怀揣着传播单位价值、展示团队风采的理想,一头扎进了稿件撰写的海洋。我的目标很简单,就是通过文字的力量,让外界听到我们的声音,感受到我们的活力。然而,理想很丰满,现实却给我上了生动的一课。 起初,我遵循传统路径,选择了一家…

QT安装和配置[安装注意点][QT找不到python27.dll][缩小空间]

安装注意点 本文摘录于&#xff1a;https://blog.csdn.net/Python_0011/article/details/131699443只是做学习备份之用&#xff0c;绝无抄袭之意&#xff0c;有疑惑请联系本人&#xff01; 双击"qt-online-installer-windows-x64-4.8.0.exe"文件后输入账号选择如下控…

什么是Capto刀柄,一起来认识一下

大家好&#xff0c;今天咱们不聊齿轮&#xff0c;说一说一款刀柄的相关内容。目前&#xff0c;高速加工中心的主轴转速可以达到10,000——50,000r /min &#xff0c;极大地提高了生产率。高速加工工具系统的主要作用是保证刀具在机床主轴中的精确定位&#xff0c;将主轴的运动和…

Python脚本启动应用并输入账号或密码

一、简介 如果每天要启动某个软件还要输入账号密码登录的需求的话&#xff0c;可以参考本文章&#xff1b; 二、Python环境 环境&#xff1a;Python3.11 已经在Windows电脑中配置Python环境变量&#xff0c;且配置了pipd的环境变量&#xff1b; 三、安装模块 安装所需要的…

计算机毕业设计 | SpringBoot招投标系统 任务发布网站(附源码)

1&#xff0c;绪论 在市场范围内&#xff0c;任务发布网站很受欢迎&#xff0c;有很多开发者以及其他领域的牛人&#xff0c;更倾向于选择工作时间、工作场景更自由的零工市场寻求零散单子来补贴家用。 如今市场上&#xff0c;任务发布网站鱼龙混杂&#xff0c;用户需要找一个…