结构篇| 浅析LLaMA网络架构
原创 程序猿阿三 程序猿阿三 2024年12月04日 22:22 福建
点击蓝字 关注我们 不迷路
01 前言
LLaMA（Large Language Model Meta AI）是由Meta AI 发布的一个开放且高效的大型基础语言模型。为什么突然讲这个模型，重要LLaMA 已经成为了最受欢迎的开源大语言模型之一，LLaMA 系列模型在学术界和工业界引起了广泛的 关注，对于推动大语言模型技能的开源发展做出了重要贡献。
第一，开源，去相识其内部模型具有可行性。第二，它很受欢迎，阐明在LLM界还是具有很强代表性，相识它内部结构有助于深入理解LLM发展路径。第三，浩繁 研究职员纷纷通过指令微调或继承预练习等方法来进一步扩展 LLaMA 模型的功 能和应用范围。此中，指令微调由于相对较低的盘算成本，已成为开辟定制化或专业化模型的首选方法，也因此出现了庞大的 LLaMA 家族。可以说LLaMA成为现在大部分互联网拥抱的对象，相识它，有助于拿下好offer。
总之，LLaMA 将有助于使 LLM 的利用和研究布衣化，是一个深度学习LLM好切入口。同时，LLaMA已在教育、法律、医疗等专业范畴有重要的应用场景，这对于构建大模型生态有先天的优势。
02 LLaMA架构
和GPT 系列一样，LLaMA 模型也是 Decoder-only 架构。底座也是Transformer的一种，《Transformer原理》和《概念篇| Transformer家族》已经介绍过Transformer模型，同样是基于自回归天生（Autoregressive）。自回归天生：在天生使命中，利用自回归（Autoregressive）方式，即逐个天生输出序列中的每个Token。在解码过程中，每次天生一个Token时，利用前面已天生的内容作为上下文，来帮助预测下一个Token。
2.0 Decoder-Only
当前主流的大语言模型都基于 Transformer 模型进行计划的。Transformer 是由多层的多头自留意力（Multi-head Self-attention）模块堆叠而成的神经网络模型。原 始的 Transformer 模型由编码器和解码器两个部分构成，而这两个部分实际上可以 独立利用，之前在《概念篇| Transformer家族》介绍过，例如基于编码器架构的 BERT 模型 和解码器架构的 GPT 模型 。 与 BERT 等早期的预练习语言模型相比，大语言模型的特点是利用了更长的向量 维度、更深的层数，进而包含了更大规模的模型参数，并重要利用解码器架构，对 于 Transformer 本身的结构与配置改变并不大。
与原生的Transformer的Decoder结构相比，做了以下几点改进：
Pre-normalization : 为了进步练习稳定性，LLaMA 对每个 transformer 子层的输入进行归一化，利用 RMSNorm 归一化函数，Pre-normalization 由Zhang和Sennrich引入。利用 RMSNorm 的好处是不用盘算样本的均值，速度提拔了40%。
SWiGLU:为了进步模型性能，结构上利用门控线性单位，且为了保持 FFN 层参数量不变，将隐蔽单位的数量调解为 234d 而不是 PaLM 论文中的 4d，同时将 ReLU 更换为 SiLU 激活，引入以进步性能。
Rotary Embeddings:为了更好地建模长序列数据，模型的输入不再利用 positional embeddings，而是在网络的每一层添加了positional embeddings (RoPE)，RoPE 方法由Su等人引入。
Grouped-Query Attention GQA：为了平衡效率和性能，部分版本采用了分组查询留意力机制。
固然从LLaMA1到LLaMA3已经发布多个版本，大要架构基本相似，接下来针对LLaMA改进点进行详细介绍，别的组件可以参考原先文章《Transformer原理》。
2.1 RMSNorm
在之前的Transformer我们提到过，LN是对单个数据的指定维度进行Norm处理与batch无关。Transformer中的Normalization层一般都是采用LayerNorm来对Tensor进行归一化，LayerNorm的公式如下：
而RMSNorm就是LayerNorm的变体，RMSNorm省去了求均值的过程,也没有了偏置 β ：
RMSNorm实现源码：
- class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim))
- def _norm(self, x): 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) return output * self.weight
- [/code] [size=4][b]2.2 激活函数的改进-SwiGLU[/b][/size]
- 与标准的Transformer一样，颠末Attention层之后就进行FeedForward层的处理，但LLama2的FeedForward与标准的Transformer FeedForward有一些渺小的差异：
- SwiGLU激活函数是SwiGLU是GLU的一种变体，此中包含了GLU和Swish激活函数。
- [size=3][b]2.2.1 GLU[/b][/size]
- 门控线性单位 GLU: 定义为门控线性单位（ Gated Linear Units, GLU），定义为输入的两个线性变换的逐元素乘积，此中一个颠末了 sigmoid 激活（也可以用其他激活函数更换）。
-
- [align=center][img=334,52]https://i-blog.csdnimg.cn/img_convert/945754b399dc2fe38119de6951a9d1f9.png[/img][/align]
- [size=3][b]2.2.2 FFN_SwiGLU[/b][/size]
- FFN_SwiGLU 原版实现利用 Swish 稍有差别，LLaMA 官方提供的代码利用 F.silu() 激活函函数：
-
- [align=center][img=484,54]https://i-blog.csdnimg.cn/img_convert/d260fb9f9e0b7cbff810059cebdb7c09.png[/img][/align]
-
- [list]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [*]
- [/list] [code]class FeedForward(nn.Module): def __init__( self, dim: int, hidden_dim: int, multiple_of: int, ffn_dim_multiplier: Optional[float], ): super().__init__() hidden_dim = int(2 * hidden_dim / 3) # custom dim factor multiplier 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))
- [/code] [size=3][b]2.3 RoPE 旋转位置编码[/b][/size]
- 必须利用位置编码，是因为纯粹的 Attention 模块是无法捕捉输入次序的，即无法理解差别位置的 token 代表的意义差别。认识《》都知道，在做留意力机制的时间，并没有考虑词语的次序，但是NLP不用文字次序会影响文本的根本性意思。比如，输入文本为“我爱吃肉包”或“肉包爱吃我”，模型会将这两句话视为相同的内容，因为嵌入中并没有明白的次序信息让模型去学习。
- [size=3][b]2.3.1 绝对编码与相对编码[/b][/size]
- 在标准的Transformer中通常是在整个网络进入Transformer Block之前做一个位置编码：
-
- [align=center][img=970,254]https://i-blog.csdnimg.cn/img_convert/17986f0920843024aa04eaeccdd8a64b.png[/img][/align]
- Transformer论文中，利用正余弦函数表示绝对位置，通过两者乘积得到相对位置。因为正余弦函数具有周期性，可以很好地表示序列中单词的相对位置。比较经典的位置编码用公式表达就是：
-
- [align=center][img=417,168]https://i-blog.csdnimg.cn/img_convert/0c08d2107c5939373061a006439bb9b7.png[/img][/align]
- 此中，i表示token在序列中的位置，设句子长度为 L,则i=0，.......,L-1 。p是token的位置向量，p（i,2t）表示这个位置向量里的第t个元素，t表示奇数维度，2t表示偶数维度；d表示token的维度。
- [size=3]除了绝对编码，还有一种相对编码，相对位置编码是根据单词之间的相对位置关系来盘算位置编码。这种编码方式更加灵活，能够捕捉到差别单词之间的相对位置信息，有助于模型更好地理解序列中单词之间的关系。但是也有缺点，盘算效率低下，同时大部分相对编码都没有落地可行性。[/size]
- RoPE（Rotary Position Embedding）旋转位置编码，由模型 RoFormer: Enhanced Transformer with Rotary Position Embedding 提出。RoPE 的核心思想是将位置编码与词向量通过旋转矩阵相乘，使得词向量不光包含词汇的语义信息，还融入了位置信息，其具有以下长处：
- [b]相对位置感知[/b]：利用绝对位置编码来到达相对位置编码的结果，RoPE 能够自然地捕捉词汇之间的相对位置关系。
- [b]无需额外的盘算[/b]：位置编码与词向量的结合在盘算上是高效的。
- [b]适应差别长度的序列[/b]：RoPE 可以灵活处理差别长度的输入序列。
- 详细如何做到呢？，这里面涉及了很多数学上推理，大家可以看一下https://spaces.ac.cn/archives/8130，我这里只做个简朴介绍一下：RoPE 借助了复数的思想，出发点是通过绝对位置编码的方式实现相对位置编码。根据复数乘法的几何意义，上述变换实际上是对应向量旋转，以是位置向量称为“旋转式位置编 码”。本质还是利用绝对位置编码，通过内积方式，得到相对位置表达式。
-
- [align=center][img=431,68]https://i-blog.csdnimg.cn/img_convert/13ba4244e34e8c7e25676e644e3b2f37.png[/img][/align]
- 根据内积满足线性叠加的性子，恣意偶数维的 RoPE，都可以表示为二维情形的拼接，即：
-
- [align=center][img=723,257]https://i-blog.csdnimg.cn/img_convert/0b889b9bdb85c54e3b012d399754f48e.png[/img][/align]
- 如果放在二维空间维度来看，可以用极坐标来理解, 旋转角度不影响轴长度：
-
- [align=center][img=1080,619]https://i-blog.csdnimg.cn/img_convert/32e0aff500b500f88cbde85ea87135b3.png[/img][/align]
-
- [list]
- [*]
- [*]
- [/list] [code]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, dtype=torch.float32) freqs = torch.outer(t, freqs) freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 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)) 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 模块利用 xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
- [/code] [size=4][b]2.4 GQA[/b][/size]
- 在说这个之前，我们先回首一下Transformer的多头留意力机制。
-
- [align=center][img=720,580]https://i-blog.csdnimg.cn/img_convert/d725d59dbf7a3815f72c6b3c16b93a89.png[/img][/align]
- 在 Transformer 中，留意力模块会并行多次重复盘算。每个并行盘算称为一个留意力头（Attention Head）。留意力模块将其查询 Query 、键 Key和值 Value的参数矩阵进行 N 次拆分，并将每次拆分独立通过一个单独的留意力头。末了，所有这些相同的留意力盘算会合并在一起，产生终极的留意力分数。能够更细致地捕捉并表达每个词汇之间的多种接洽和微妙差异。
-
- [align=center][img=1080,374]https://i-blog.csdnimg.cn/img_convert/bc7eed132241cc571f305f415831024d.png[/img][/align]
-
- [list]
- [*] 在 [b]MHA（Multi Head Attention）[/b] 中，每个头有本身单独的 key-value 对；标准的多头留意力机制，h个Query、Key 和 Value 矩阵。
- [*] 在 [b]MQA（Multi [/b][b]Query[/b][b] Attention）[/b] 中只会有一组 key-value 对；多查询留意力的一种变体，也是用于自回归解码的一种留意力机制。与MHA差别的是，MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。
- [*] 在 [b]GQA（Grouped [/b][b]Query[/b][b] Attention）[/b]中，会对 attention 进行分组操作，query 被分为 N 组，每个组共享一个 Key 和 Value 矩阵GQA将查询头分成G组，每个组共享一个Key 和 Value 矩阵。GQA-G是指具有G组的grouped-query attention。GQA-1具有单个组，因此具有单个Key 和 Value，等效于MQA。而GQA-H具有与头数相当的组，等效于MHA。
- [/list] GQA介于MHA和MQA之间。GQA 综合 MHA 和 MQA ，既不损失太多性能，又能利用 MQA 的推理加快。不是所有 Q 头共享一组 KV，而是分组肯定头数 Q 共享一组 KV，比如上图中就是两组 Q 共享一组 KV。现在LLaMA3基本都利用了GQA结构。
-
- [list]
- [*]
- [*]
- [*]
- [/list] [code]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): 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], ): bsz, seqlen, _ = x.shape # 计算 q、k 、v 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) # 加上Rope位置旋转 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) # kv 缓存 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] # GQA应用 # repeat k/v heads if n_kv_heads < n_heads keys = repeat_kv( keys, self.n_rep ) # (bs, cache_len + seqlen, n_local_heads, head_dim) values = repeat_kv( values, self.n_rep ) # (bs, cache_len + seqlen, n_local_heads, head_dim)
- xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim) keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim) values = values.transpose( 1, 2 ) # (bs, n_local_heads, cache_len + seqlen, head_dim) 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)
Python的完整的LLaMa3代码在github可以快速找到，其核心代码也不过几百行，但此中的计划思想和理念，够我们这些小白喝一段时间，盼望通过不停深入学习，进步对LLM实际的理解。通过记录所学的知识，建立本身的体系性思维。
感谢您完成阅读
参考文章:
https://mingchao.wang/BDL2PIY6/
https://github.com/meta-llama/llama3/blob/main/llama/model.py
https://jalammar.github.io/illustrated-transformer/
https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/MODEL_CARD.md
https://github.com/harleyszhang/llm_note/blob/main/1-transformer_model/llama1-3%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E8%AF%A6%E8%A7%A3.md
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
