为什么对 LLaMA 系列模型的关键技术进行总结?

  • LLaMA 系列模型贡献或证实了很多有用的技术点,这些点推动了大模型领域的发展。
  • 本文主要介绍:RMS normalization、GQA、SwiGLU、RoPE 这 4 个技术点。
  • 对于每一个技术点,会介绍其研究背景、主要内容、代码实现这 3 个方面。

1、RMS normalization

(1)提出背景

归一化方法,包括 batch normalization、layer normalization、RMS norm,这些方法核心作用,都是为了解决模型在训练过程中-模型训练难度大、模型训练收敛慢的问题。

RMS norm 是对 layer normalization 的改进,是为了进一步降低 layer norm 的计算量。

(2)主要内容

第一,RMSNorm 的全称是 root mean square normalization,RMSNorm 核心是对 Layer Normalization 的改进,去掉了 Layer Normalization 中的平移项,包括均值项和可学习的 bias 项,RMS Norm 在LLaMA中位置如下右图所示。

Transformer versus LLaMA

第二,精度上 RMS Norm 和 Layer Norm 持平,但计算上效率更高,当前主流的 LLM 通常都采用 RMS Norm 方法。

第三,RMSNorm 和 Pre-Norm 是两个不同的概念。前者是一种具体的 normalization 方法(例如 layer norm 或 RMS norm),后者通常是指 normalization 在 transformer block 中放置的位置。

Pre-Norm 是 normalization 模块放在 Attention 或 FFN 的前面,Post-Norm 是 normalization 放在 Attention 或 FFN 的后面。

从 GPT2 开始,大家在训练大参数量模型时更多使用 Pre-Norm,因为在模型参数量增加时 Post-Norm 很难训练,容易导致梯度消失。

(3)Layer Norm 的 Pytorch 代码

class LayerNorm(nn.Module):    def __init__(self, features, eps=1e-5):        super().__init__()    self.eps = eps    self.alpha = nn.Parameter(torch.ones(features))    self.bias = nn.Parameter(torch.zeros(features))     def forward(self, x):        x_mean = x.mean(dim=-1, keepdim=True)        x_var = x.var(dim=-1, keepdim= True)    return self.alpha*(x-x_mean)/torch.sqrt(x_var+self.eps)+self.bias

(4)RMS Norm 的 Pytorch 代码

class RMSNorm(nn.Module):    def __init__(self, features, eps=1e-5):      super().__init__()    self.eps = eps    self.alpha = nn.Parameter(torch.ones(features))  def forward(self, x):    x_var = x.pow(2).mean(dim=-1, keepdim= True)    return self.alpha*x/torch.sqrt(x_var+self.eps)

2、GQA(Group Query Attention)

(1)提出背景

对于 multi-head attention(MHA),人们提出了多种优化方案:

sparse attention 通过在 QK 之间采用稀疏计算降低了整体的计算量(包含 train 和 test);

flash attention 通过分块和重计算的方案减少了 MHA 在计算(包含 train 和推理)过程中的存储访问开销;

在推理过程中,KV cache 通过 cache 住 MHA 计算过程中的 KV,虽然内存消耗线性增长,但可以极大降低矩阵计算,提升推理速度;

在推理过程中,KV cavhe 本质上是使用空间换时间,那么如何降低 KV cache 的空间呢?

multi query attention(MQA)通过共用KV降低了在推理过程中的 KV cache 消耗,加快了推理速度但精度会降低;

在推理过程中,能不能提出一种方案,能够在推理过程中加快速度同时保持 MHA 的精度呢?

也就有了 group query attention(GQA),是 MHA 和 MQA 的一种折中方案,GQA 和 MQA 本质上都是对 KV cache 的优化。

(2)主要内容

第一,MHA、MQA、GQA 的核心不同在于 query、key、value 三者 head 个数的不同,如下图所示**(每一个小块代表是 heads 个数)。**

MHA、MQA、GQA 的示意图

MHA:head 个数 query=key=value=N

MQA:head 个数 query=N,key=value=1

GQA:head 个数 query=N,key=value=N/group

GQA 通过分组的方式,减少了需要缓存的键和值的 cache,从而减少了内存的使用。

第二,在具体实现上,因为 attention 计算需要 QKV 三者的 dimension 维度相同,那么 GQA 如何实现呢?

在训练过程中,通过将 KV 的维度"复制" group 份,来实现 QKV 的计算;在推理过程中,多个 heads 的 Q 共用一组 KV 实现 KV cache 内存降低。

GQA 中的 QKV 的获取方式

(3)MHA 的代码实现

class Multiheadattention(nn.Module):    def __init__(self, config):        super().__init__()    self.hidden_size = config.hidden_size    self.wq = nn.Linear(config.hidden_size,config.hidden_size)    self.wk = nn.Linear(config.hidden_size,config.hidden_size)    self.wv = nn.Linear(config.hidden_size,config.hidden_size)    self.wo = nn.Linear(config.hidden_size,config.hidden_size)    self.num_heads = config.num_heads    self.d_perhead = config.d_perhead    self.attention = Attention(config)    def forward(self, x, attention_mask):    batch_size, seq_length = x.shape[0], x.shape[1]    Q = self.wq(x)    K = self.wk(x)    V = self.wv(x)    Q = Q.reshape(batch_size,seq_length,self.num_heads,self.d_perhead).transpose(1, 2)    K = K.reshape(batch_size,seq_length,self.num_heads, self.d_perhead).transpose(1, 2)    V = V.reshape(batch_size,seq_length,self.num_heads, self.d_perhead).transpose(1, 2)    temp = Attention(Q, K, V, attention_mask)    temp = temp.transpose(1,2).reshape(batch_size, seq_length, self.hidden_size)    output = self.wo(V)    return output

(4)GQA 的代码实现

def repeat_kv(KV, num):    batch, num_kv_heads, sequence_len, head_dim = KV.shape    return KV.reshape(batch,num_kv_heads, 1, sequence_len, head_dim).expand(batch,num_kv_heads, num, sequence_len, head_dim].reshape(batch,num_kv_heads*num, sequence_len, head_dim)class GroupQA(nn.Module):    def __init__(self, config):        super().__init__()    self.hidden_size = config.hidden_size        self.num_q_heads = config.num_q_heads        self.num_kv_heads = config.num_kv_heads        self.d_perhead = config.d_perhead        self.attention = Attention(config)    self.wq = nn.Linear(config.hidden_size,config.hidden_size)    self.wk = nn.Linear(config.hidden_size,self.num_kv_heads * self.d_perhead)    self.wv = nn.Linear(config.hidden_size,self.num_kv_heads * self.d_perhead)    self.wo = nn.Linear(config.hidden_size,config.hidden_size)          def forward(self, x, attention_mask):      batch_size, seq_length = x.shape[0], x.shape[1]    Q = self.wq(x)    K = self.wk(x)    V = self.wv(x)    Q = Q.reshape(batch_size,seq_length,self.num_q_heads,self.d_perhead).transpose(1,2)    K = K.reshape(batch_size,seq_length,self.num_kv_heads, self.d_perhead).transpose(1, 2)    V = V.reshape(batch_size,seq_length, self.num_kv_heads, self.d_perhead).transpose(1, 2)        K = repeat_kv(K, self.num_q_heads// self.num_kv_heads)        V = repeat_kv(V, self.num_q_heads// self.num_kv_heads)    temp = Attention(Q, K, V, attention_mask)    temp = temp.transpose(1,2).reshape(batch_size, seq_length, self.hidden_size)    output = self.wo(V)    return output

3、SwiGLU

(1)提出背景

激活函数是神经网络模型的重要组成,核心作用是为网络**“提供非线性”,同时在反向传播时能够确保有“稳定的梯度流”**。

大模型中常用的激活函数包括:ReLU、sigmoid、Swish、GLU,定义如下图所示。

目前的激活函数仍有一些待完善的地方:

  • ReLU:梯度在负区间为 0,导致神经元死亡;
  • sigmoid:输入数值在远离 0 时会存在梯度消失,靠近 0 时会存在梯度爆炸;
  • Swish:缺乏门控机制;
  • GLU:使用 sigmoid 作为门控函数,容易导致梯度消失或梯度爆炸。

不同种类激活函数的定义

不同种类激活函数的数值

(2)主要内容

第一,SwiGLU 是一种结合 Swish 和 GLU 机制的激活函数,其是 GLU 的变体并使用 Swish 作为其门控部分,定义如下图所示。

SwiGLU 激活函数的定义

第二,需要说明的是:通常在 SwiGLU 激活函数最后加一个线性层,来保证维度不变,如下图所示。

Original FFN versus SwiGLU+Linear

第三,在效果方面,SwiGLU 需要进行三次矩阵乘法运算,相比于 ReLU 等激活函数计算复杂度更高,但能实现更好的梯度回传、泛化性、效果。

(3)代码实现

class SwiGLU(nn.Module):    def __init__(self, config):      super().__init__()    self.dim = config.dim    self.hidden_dim = config.hidden_dim    self.w1 = nn.Linear(config.dim, config.hidden_dim, bias = False)    self.w2 = nn.Linear(config.dim, config.hidden_dim, bias = False)    self.w3 = nn.Linear(config.hidden_dim, config.dim, bias = False)   # nn.Linear默认情况下bias为true。    def forward(self, x):    x1, x2 = x.chunk(2, dim=-1)    return self.w3(F.silu(self.w1(x1))*self.w2(x2))

4、RoPE

(1)提出背景

Attention 机制本身不包含位置信息,所以有了 Transformer 论文中的 Positional encoding,但原始 Transformer 论文中的 Positional encoding 存在泛化性不足的问题。

Attention 在计算两个向量的内积时,只会关心两个向量之间的相对位置,并不会关心两个向量的绝对位置,这是 Attention 计算原理决定的。那么该如何更好的表示 QK 两个向量之间的相对距离呢?

(2)主要内容

第一,旋转位置编码 RoPE 作为一种相对位置编码,核心是通过由正弦和余弦组成的旋转矩阵,来对向量进行旋转,由此构建 QK 向量之间的相对位置关系。

两个向量的相对位置可以使用旋转矩阵来进行表示,语义相似性使用向量长度表示,位置相似性就是旋转的角度大小。

第二,在做 Attention 计算时,会先进行多 head 投影,再加上旋转位置编码,并不是像原始 Transformer 论文中在 embedding 之后紧接着 positional encoding。此外,只对 QK 做位置编码,对 V 不做。

第三,旋转位置矩阵的最终形式如下图所示:其中 x 为输入向量 Q 或 K,R 为对应的旋转矩阵,d 为 embedding 维度,m 代表绝对位置,Hadamard 代表逐个元素计算。

RoPE 最终计算公式

(3)RoPE 原理证明

原理整理主要包括—什么是二维旋转矩阵、二维旋转矩阵特性、QK 计算与旋转矩阵、如何从二维旋转矩阵扩充到多维、最终的计算这 5 个部分。

第一,什么是二维旋转矩阵?经过二维旋转矩阵之后,向量大小不变,改变了旋转角度,如下图所示。

第二,二维旋转矩阵两个特性?两个旋转矩阵相乘=旋转矩阵角度相加;旋转矩阵转置=旋转矩阵负角度。

第三,QK 计算与旋转矩阵有什么关系?旋转矩阵引入 QK 计算引入了相对位置的信息。

第四,如何从二维旋转矩阵扩充到多维?从二维依次填充到多维。

第五,最终的计算公式。

Computational efficient realization of RoPE

(4)RoPE 的代码实现

def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):    # 计算词向量元素两两分组之后,每组元素对应的旋转角度\theta_i    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))    # 生成 token 序列索引 t = [0, 1,..., seq_len-1]    t = torch.arange(seq_len, device=freqs.device)    # freqs.shape = [seq_len, dim // 2]     freqs = torch.outer(t, freqs).float()  # 计算m * \theta    # 计算结果是个复数向量    # 假设 freqs = [x, y]    # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)     return freqs_cis# 旋转位置编码计算def apply_rotary_emb(xq,xk,freqs_cis) :    # xq.shape = [batch_size, seq_len, dim]    # xq_.shape = [batch_size, seq_len, dim // 2, 2]    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)    # 转为复数域    xq_ = torch.view_as_complex(xq_)    xk_ = torch.view_as_complex(xk_)    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2)   # xq_out.shape = [batch_size, seq_len, dim]    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)    return xq_out.type_as(xq), xk_out.type_as(xk)class Attention(nn.Module):    def __init__(self, args: ModelArgs):        super().__init__()        self.wq = Linear(...)        self.wk = Linear(...)        self.wv = Linear(...)        self.freqs_cis = precompute_freqs_cis(dim, max_seq_len * 2)    def forward(self, x: torch.Tensor):        bsz, seqlen, _ = x.shape        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)        xq = xq.view(batch_size, seq_len, dim)        xk = xk.view(batch_size, seq_len, dim)        xv = xv.view(batch_size, seq_len, dim)        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)        scores = torch.matmul(xq, xk.transpose(1, 2)) / math.sqrt(dim)        scores = F.softmax(scores.float(), dim=-1)        output = torch.matmul(scores, xv)  # (batch_size, seq_len, dim)        return output

最后

为什么要学AI大模型

当下,⼈⼯智能市场迎来了爆发期,并逐渐进⼊以⼈⼯通⽤智能(AGI)为主导的新时代。企业纷纷官宣“ AI+ ”战略,为新兴技术⼈才创造丰富的就业机会,⼈才缺⼝将达 400 万!

DeepSeek问世以来,生成式AI和大模型技术爆发式增长,让很多岗位重新成了炙手可热的新星,岗位薪资远超很多后端岗位,在程序员中稳居前列。

在这里插入图片描述

与此同时AI与各行各业深度融合,飞速发展,成为炙手可热的新风口,企业非常需要了解AI、懂AI、会用AI的员工,纷纷开出高薪招聘AI大模型相关岗位。
在这里插入图片描述
最近很多程序员朋友都已经学习或者准备学习 AI 大模型,后台也经常会有小伙伴咨询学习路线和学习资料,我特别拜托北京清华大学学士和美国加州理工学院博士学位的鲁为民老师给大家这里给大家准备了一份涵盖了AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频 全系列的学习资料,这些学习资料不仅深入浅出,而且非常实用,让大家系统而高效地掌握AI大模型的各个知识点。

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

AI大模型系统学习路线

在面对AI大模型开发领域的复杂与深入,精准学习显得尤为重要。一份系统的技术路线图,不仅能够帮助开发者清晰地了解从入门到精通所需掌握的知识点,还能提供一条高效、有序的学习路径。

img

但知道是一回事,做又是另一回事,初学者最常遇到的问题主要是理论知识缺乏、资源和工具的限制、模型理解和调试的复杂性,在这基础上,找到高质量的学习资源,不浪费时间、不走弯路,又是重中之重。

AI大模型入门到实战的视频教程+项目包

看视频学习是一种高效、直观、灵活且富有吸引力的学习方式,可以更直观地展示过程,能有效提升学习兴趣和理解力,是现在获取知识的重要途径

在这里插入图片描述
光学理论是没用的,要学会跟着一起敲,要动手实操,才能将自己的所学运用到实际当中去,这时候可以搞点实战案例来学习。
在这里插入图片描述

海量AI大模型必读的经典书籍(PDF)

阅读AI大模型经典书籍可以帮助读者提高技术水平,开拓视野,掌握核心技术,提高解决问题的能力,同时也可以借鉴他人的经验。对于想要深入学习AI大模型开发的读者来说,阅读经典书籍是非常有必要的。
在这里插入图片描述

600+AI大模型报告(实时更新)

这套包含640份报告的合集,涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师,还是对AI大模型感兴趣的爱好者,这套报告合集都将为您提供宝贵的信息和启示。
在这里插入图片描述

AI大模型面试真题+答案解析

我们学习AI大模型必然是想找到高薪的工作,下面这些面试题都是总结当前最新、最热、最高频的面试题,并且每道题都有详细的答案,面试前刷完这套面试题资料,小小offer,不在话下
在这里插入图片描述

在这里插入图片描述

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

Logo

码道开发者社区,聚焦华为云码道 CodeArts 代码智能体,沉淀 Agent、Skill、鸿蒙开发实战内容,供开发者查阅资料、交流技术、分享工程实践

更多推荐