结构篇| 浅析LLaMA网络架构 - ToB企服应用市场:ToB评测及商务社交产业平台

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)

复制代码