开门见山,本篇文章主要包含以下知识点:
- 什么是Transformer
- 大模型的结构设计
- 模型微调
- 分布式训练
- 大模型量化
如果能根据这个分布模块每个部分都能讲个大概,大抵就是掌握大模型的训练技巧了。
参考文献:> https://blog.csdn.net/gitblog_00158/article/details/141121324
1. What is Transformer
知识点大纲:
- 预备知识
- Transformer的框架结构
1.1 预备知识
引用参考:> https://blog.csdn.net/weixin_42475060/article/details/121101749
Transformer本质是一个神经网络,那他必然避不开很多底层设计的知识点,比如:升/降维,归一化,激活函数,归纳偏置,反向梯度传播,优化器,显存占用计算等等。在这一节我们将统一介绍。
Transformer的关键组件就是Attention模块与全连接层模块FFN设计,了解这两个设计是了解整个Transformer的前提。
什么是注意力?简单来说,就是给定一个序列,为了完成任务,我该更关注这个序列有哪些结构,将独立的序列信息进行聚合,从而得到序列中的关键信息。那么这种进行关注的方法具体该如何实现,这就要引入注意力矩阵了。
我们知道,相关性一般是通过点积来得到的。第一步,对于一个序列来说,自注意力模块会生成对应的三个向量,即查询向量、键向量和值向量,三个向量分别由三个可学习的权重矩阵与嵌入向量进行矩阵乘法得到,一般来说,新向量在维度上往往比嵌入向量更低。为什么需要三个不同的向量来和嵌入向量进行交互呢?准确来说,QKV更多是检索融合思想的具象计算方式,首先,查询向量与键向量的相似度计算方式就是通过归一化点积实现,然后通过缩放后加权求和对应的值向量即可得到提取信息的表达。另外,QKV构成了信息检索的最小完备集,在数学上来说,增加向量是冗余的,更多的表达与结构复杂性往往是通过多头注意力机制实现的。

计算自注意力的第二步就是计算得分,也就是我们熟悉的上图公式,一般来说,下图的公式会涉及到以下的几个问题:
(1)Q: 什么是Softmax?有什么样的应用场景?
A: 将输入张量的值转换为概率分布(所有值在[0,1]之间,且和为1)。适用于 NLP(Transformer 注意力)、分类任务等。dim 指定在哪个维度计算 softmax(默认 -1 表示最后一个维度)。SafeSoftmax的设计是先减去最大值(x_max)防止数值溢出(exp 可能导致 inf)。这样的好处主要体现在大梯度输入与混合精度训练中防止数值溢出(梯度与FP16)与数值差异过大(注意力分数导致输出概率错误)的问题。
点击查看代码
import torch
def softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
# 防止数值溢出,减去最大值
x_max = x.max(dim=dim, keepdim=True).values
x_exp = torch.exp(x - x_max)
return x_exp / x_exp.sum(dim=dim, keepdim=True)
(2)Q: \(d_k\)是什么?为什么要除以它?
A: \(d_k\)是模型结构设计时设定的超参数,防止点积结果过大导致梯度饱和。一般来说,嵌入矩阵X的形状是\(n*d_{model}\),注意力计算中的各学习矩阵\(W_q\),\(W_k\),\(W_v\)的形状会设定为\(d_{model}*d_k\),\(d_{model}*d_k\),\(d_{model}*d_v\)。所以本质来看,d_k就是Q与K的维度超参数,我们不希望过大的点积结果和过小的点积结果在经过softmax后会导致异常值被exp()放大到无穷大,其他值被压缩到接近0,通过缩放可以将点积结果的方差稳定在1左右,避免了梯度消失,让模型训练过程更加稳定和高效。
(3)Q: Self-Attention的其他表达方法?
A:Self-Attention 的核心目标是 建模序列内部元素之间的相关性。只要能达到这个目的,并且满足一些工程上的期望,其他的形式也是可能的。理想的 Attention 机制应该具备:
高效计算: 点积可以通过高度优化的矩阵乘法实现,计算速度块;
强大的表达能力: Query 可以主动去“查询”并关注到相关的 Key,然后根据相关性在 Value 上提取信息,同时忽略不相关的部分。
足够的模型容量: 引入\(W_q, W_k, W_v\)线性映射:输入x并不是直接作为 Q, K, V,而是先经过三个不同的线性变换(乘以权重矩阵\(W_q, W_k, W_v\))得到 Q, K, V。这增加了模型的参数量和学习能力,让模型可以学习到在不同“子空间”中进行查询、匹配和信息提取。多头注意力,输出线性映射。Scaled Dot-Product Attention在效率、表达能力和模型容量之间取得了很好的平衡。
自注意力计算完成后,我们进一步简单介绍多头注意力机制,其本质就是用多个单头注意力得到注意力计算结果,拼接成一个注意力矩阵再与附加权重矩阵相乘换维,然后送入FFN。
1、扩展了模型专注于不同位置的能力。
2、有多个查询/键/值权重矩阵集合,(Transformer使用八个注意力头)并且每一个都是随机初始化的。和上边一样,用矩阵X乘以\(W_Q\),\(W_K\),\(W_V\)来产生查询、键、值矩阵。
3、self-attention只是使用了一组\(W_Q\),\(W_K\),\(W_V\)来进行变换得到查询、键、值矩阵,而Multi-Head Attention使用多组\(W_Q\),\(W_K\),\(W_V\)得到多组查询、键、值矩阵,然后每组分别计算得到一个Z矩阵。特别地,除了第0个编码器,其他编码器都不需要进行词嵌入,它可以直接将前一层的编码器输出作为输入。
因此,
关注不同子空间的信息: 单个 Attention 头可能只能关注到一种类型的相关性(比如,主要关注句法依赖)。多头允许模型同时从不同的表示子空间学习信息。你可以类比 CNN 中的不同卷积核(Channel),它们会学习提取图像的不同特征(边缘、纹理、颜色等)。不同的 Attention 头也类似,可能有的头关注短距离依赖,有的关注长距离依赖,有的关注特定词性关系等。
增强模型表达能力(Again!): 多头实际上是多个 Attention 机制的集成(ensemble),每个头可以捕捉不同的上下文依赖关系。这比单个大头具有更强的表达能力,可以更细致地建模复杂的依赖关系。
稳定学习过程 (潜在好处): 每个头在较低维度上操作,可能比在原始高维空间直接计算 Attention 更稳定。
并行计算: 各个头的计算是独立的,可以并行执行,计算效率高。
位置编码也是Transformer中的一个重要组成,如果不添加位置编码,那么无论单词在什么位置,它的注意力分数都是确定的。这不是我们想要的。为了理解单词顺序,Transformer为每个输入的词嵌入添加了一个向量,这样能够更好的表达词与词之间的关系。词嵌入与位置编码相加,而不是拼接,他们的效率差不多,但是拼接的话维度会变大,所以不考虑。
sin/cos pos embed:从二进制编码到有界连续的周期函数编码,1/10000作为频率的一个附值主要是将函数的波长拉长,避免不同t的情况下两个token位置过近,单独的sin可以实现绝对位置编码,那么通过sin与cos的线性转换,就可以得到相对位置编码。这也就是旋转位置编码,和绝对位置编码相比,RoPE具有更好的外推性,目前是大模型相对位置编码中应用最广的方式之一。在这之后对于RoPE还有一个优化,通过线性项\(-abs{m(i-j)}\)来抵消标记i,j之间的所有注意力logits并变为\(q^T_ik_j-|m(i-j)|\)。为此,MPT-7B 代码将偏移矩阵实现为注意力logits中的附加项。为了使用 LM-Infinite 进行增强,我们只需将偏移矩阵裁剪为最小值\(-abs{mL_{pretrain}}\)
理论推导可以参考:> https://cloud.tencent.com/developer/article/2327751
经过多头注意力机制后,还需要进行Add&Normalize的处理。简单来说,就是残差归一化。
加入残差块的目的是为了防止在深度神经网络的训练过程中发生退化的问题,退化的意思就是深度神经网络通过增加网络的层数,Loss逐渐减小,然后趋于稳定达到饱和,然后再继续增加网络层数,Loss反而增大。
Normalize归一化目的:
这样做的好处是:
1.让数值进入激活函数的敏感区间: 类似于前面提到的 softmax,很多激活函数(如 sigmoid, tanh)在输入值过大或过小时,梯度会趋近于 0。Norm 操作让数据落在梯度较大的区域。
2.缓解梯度消失/爆炸: 稳定的数据分布有助于梯度的稳定传播。
3.降低对初始化参数的敏感度: 让模型更容易训练。
4.起到一定的正则化效果。
LN是在同一个样本中不同神经元之间进行归一化,而BN是在同一个batch中不同样本之间的同一位置的神经元之间进行归一化。
LayerNorm旨在缓解ICS问题,它对 单个样本 的 所有特征 进行归一化,使其均值为0,方差为1。
$ LayerNorm(x)=\gamma \cdot {\frac{x-\mu_L}{\sqrt{\sigma^2_L+\epsilon}+\beta $
BN是对于相同的维度进行归一化,但是NLP中输入的都是词向量,一个300维的词向量,单独去分析序列数据的每一维也没有意义,因此这里选用的是LN。
那么RMSNorm是什么?简单来说,RMSNorm就是移除了均值中心化步骤的LayerNorm。
优势:
计算效率高:相比LayerNorm,减少了约7%-64%的计算时间(根据论文数据)。移除了均值计算这一涉及多步操作和同步的步骤。
内存占用低:无需存储均值相关中间变量。
性能相当:大量实验表明,在保持与LayerNorm相当的模型性能的同时实现了加速。
为何可行?
Pre-Norm的铺垫:在Pre-Norm结构中,归一化的主要目的是控制尺度而非严格对齐分布。移除均值中心化对模型性能影响不大。实际发现,如果输入数据的均值发生了变化,但数据的分布形状和范围保持不变,那么具有 recentering invariance 的算法或函数的输出应该不受影响。
硬件优化需求:对于大规模并行计算(如GPU),减少跨通道/时间步的依赖(如均值计算)能显著提升硬件利用率。
与RoPE等位置编码兼容:一些研究认为,保留输入的原始均值信息可能对某些类型的位置编码(如RoPE)更有利。
另外,归一化的位置也是有所不同的,早期BERT采用的是Post-Norm + LayerNorm,到现代大语言模型青睐的PreNorm+RMSNorm,那么这些的区别在哪里呢?
1.梯度消失:梯度信号是从输出层逐层传递回输入层来更新参数的,根据链式法则,梯度计算涉及多个雅可比矩阵的连乘,如果单层变换的梯度持续小于1,那么回传过程会呈指数级衰减,导致靠近输入层的参数几乎无法更新。
2.内部协变量偏移(ICS):ICS是指在训练过程中,由于前一层参数的不断更新,导致后续网络层的输入数据分布持续发生变化的现象。这迫使网络层不断适应新的输入分布,而非专注于学习有效的特征表示,从而降低了训练效率,并可能导致梯度不稳定,需要更小的学习率和更精细的初始化。
我们知道残差链接可以通过引入捷径来缓解梯度消失,PostNorm的思想就是将归一化放在残差链接之后,这样的优势在于:
理论性能上限高:归一化作用于最终输出,保证了每层输出的特征尺度一致,有利于直接抽取中间层特征(如BERT)。它维持了网络的“真实”深度,同等层数下理论表征能力可能更强。
微调友好:其固有的梯度衰减倾向(见下文)在微调时可能成为优势,因为它天然抑制了底层预训练参数的剧烈变动,有助于保留预训练知识。
正则化效果强:对残差块的完整输出进行归一化,正则化效果更全面。
但问题在于:
训练不稳定与梯度衰减放大
初始化敏感:在训练初期,\(x_t\)和\(F(x_t)\)的方差叠加可能导致\(x_t + F(x_t)\)的方差增大。LayerNorm会将其强制缩放回单位方差附近,例如,若初始方差为1,叠加后为2,LayerNorm近似将其乘以\(1/\sqrt{2}\)。
梯度路径受阻:这种缩放效应作用在残差主干路径上。经过 L 层后,来自浅层的原始信号\(x_0\)对\(x_L\)的贡献可能被缩减约\(1/\sqrt{2})^L\)倍,严重削弱了残差连接缓解梯度消失的能力。这使得Post-Norm模型训练困难,通常需要精心设计的学习率Warmup和较小的初始化值(如BERT的\(\mathcal{N}(0, 0.02^2)\))来避免早期梯度爆炸和后期梯度消失。
Adam的缓解:值得注意的是,现代自适应优化器(如Adam)通过对梯度进行归一化(\(\hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)\)),能在一定程度上缓解梯度绝对值过小的问题,使得即使梯度信号较弱,参数也能获得有效更新。但这并未完全消除Post-Norm的训练稳定性挑战。
这里特别提一下AdamW,Adam 和 AdamW 的核心区别在于权重衰减(Weight Decay)的处理方式,简单来说就是"耦合"与"解耦"的差异。
第一,机制不同。在标准 Adam 中,L2 正则化是加在损失函数里的,这意味着权重衰减项会被自适应学习率缩放。具体来说,Adam 的更新步长会除以梯度平方根的移动平均,导致权重衰减项也被这个因子除了一下。对于梯度大的参数,衰减效果会被意外削弱,导致正则化不一致。而 AdamW 把权重衰减从梯度更新中独立出来,直接作用于参数本身,不受自适应学习率影响,实现了真正的解耦。
第二,效果不同。AdamW 解决了 Adam 在特定场景下正则化失效的问题,泛化能力更强。特别是在 Transformer 架构(如 BERT、GPT)中,AdamW 几乎是标准配置,收敛更稳定,最终效果通常优于 Adam。论文实验表明,在同样的超参数设置下,AdamW 的测试误差通常更低。
第三,使用建议。现在训练大模型或微调任务时,默认首选 AdamW。除非是某些特定 legacy 项目或复现旧论文要求,否则不再使用带 L2 正则的 Adam。因为在实际工程中,AdamW 不仅效果更稳,超参数(比如 weight decay)也更直观,不需要为了补偿自适应缩放而去调整正则系数。
总的来说,AdamW 是对 Adam 的一个关键修正,核心贡献就是"解耦权重衰减"。在我的项目实践中,比如做 LLM 微调时,我们都是直接用 AdamW 配合 Cosine 学习率调度,这已经成为了一套标准最佳实践。
所以,PreNorm是当今LLM的更优选择:先归一化,在残差连接:
优势:
训练极其稳定:归一化仅作用于非线性变换\(F\)的输入,主干残差路径\(x_t\)的梯度流\(I\)不受任何缩放干扰。这使得梯度量级在层间传递时更加均衡,极大提升了训练稳定性,无需复杂的Warmup即可训练非常深的网络(千层级别)。
即插即用:训练稳定性高,对超参数和初始化不那么敏感。
局限:深度虚化 (Depth Virtualization)
等效宽度增加:递归展开Pre-Norm公式:
\(x_{t+1} = x_0 + F_0(\text{Norm}(x_0)) + F_1(\text{Norm}(x_1)) + \dots + F_t(\text{Norm}(x_t))\)
每一层的增量\(F_i(\text{Norm}(x_i))\)的量级大致相当。随着层数\(t\)增大,\(x_t\)的值主要由前面所有层的累加决定。新增一层\(F_t(\text{Norm}(x_t))\)对\(x_{t+1}\)的相对贡献变小,使得\(\text{Norm}(x_{t+1}) \approx \text{Norm}(x_t)\)。这意味着深层网络的效果越来越像一个不断加宽的浅层网络,而不是真正意义上的深度增加。
理论性能可能受限:由于深度虚化,Pre-Norm模型在同等参数量下,有效深度可能不如Post-Norm,这可能限制其理论性能上限。实践中,常通过增加FFN层的宽度(如LLaMA将FFN维度设为隐藏层维度的\(8/3 \times 2 = 5.33 \dots\)倍,实际实现通常取倍数如4或8/3)来补偿这种深度损失。
结论:对于追求极致训练稳定性和可扩展性的大模型(如GPT-3, LLaMA),Pre-Norm几乎成为必然选择。而对于层数不多、或需要利用中间层稳定特征表示的场景(如BERT),Post-Norm仍有其价值。
在工程实践中,Post-Norm通常需要更小的参数初始化标准差(BERT 0.02)来配合WarmUp,抑制初始阶段的梯度爆炸。
PreNorm虽然稳定,但为了弥补深度虚化,需要配合更宽的FFN层。
点击查看llama-decode代码
import torch
import torch.nn as nn
# 简化的RMSNorm实现
class RMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return (self.weight * hidden_states).to(input_dtype)
# 简化的LLaMA Decoder Layer
class LlamaDecoderLayer(nn.Module):
def __init__(self, config):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = YourAttentionModule(config) # 替换为实际的注意力模块
self.mlp = YourMLPModule(config) # 替换为实际的MLP模块
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(self, hidden_states, ...): # 省略其他参数如attention_mask
# --- Self Attention Block ---
residual = hidden_states
# 1. Pre-Normalization (RMSNorm)
normalized_hidden_states = self.input_layernorm(hidden_states)
# 2. Attention
attn_output = self.self_attn(normalized_hidden_states, ...)
# 3. Residual Connection
hidden_states = residual + attn_output
# --- FFN Block ---
residual = hidden_states
# 4. Pre-Normalization (RMSNorm)
normalized_hidden_states = self.post_attention_layernorm(hidden_states)
# 5. MLP
ffn_output = self.mlp(normalized_hidden_states)
# 6. Residual Connection
hidden_states = residual + ffn_output
return hidden_states
这个简化的LLaMA层清晰地展示了Pre-Norm(先Norm再Attention/MLP)和RMSNorm的应用,以及两次残差连接确保信息流畅传递的设计。
简单总结一下,Post-Norm 奠定了基础,但在深度和稳定性上面临挑战,更适合层数较少或需要稳定中间特征的场景。
Pre-Norm 通过牺牲部分理论深度换取了卓越的训练稳定性,成为训练超大规模模型的关键,但可能需要通过加宽网络等方式补偿。
LayerNorm 作为Transformer的标配归一化算法,有效解决了ICS问题且不受批次大小限制。
RMSNorm 作为LayerNorm的效率优化版,在Pre-Norm大行其道的背景下,凭借计算和内存优势成为新宠。

该部分参考原文:> https://blog.csdn.net/qq_54445177/article/details/147096307?spm=1001.2014.3001.5502
全连接层Feed Forward Network是一个两层的神经网络,先线性变换,然后ReLU非线性,再线性变换。
这两层网络就是为了将输入的Z映射到更加高维的空间中然后通过非线性函数ReLU进行筛选,筛选完后再变回原来的维度。
1.2 Transformer的框架结构

Transformer的编码组件是由6个编码器叠加在一起组成的,解码器同样如此。所有的编码器在结构上是相同的,但是它们之间并没有共享参数。编码器的简略结构如下:
从编码器输入的句子首先会经过一个自注意力层,这一层帮助编码器在对每个单词编码的时候时刻关注句子的其它单词。解码器中的解码注意力层的作用是关注输入句子的相关部分,类似于seq2seq的注意力。原结构中使用到的是多头注意力机制(Multi-Head Attention)。
和Encoder Block一样,Decoder也是由6个decoder堆叠而成的,Nx=6。包含两个 Multi-Head Attention 层。第一个 Multi-Head Attention 层采用了 Masked 操作。第二个 Multi-Head Attention 层的K, V矩阵使用 Encoder 的编码信息矩阵C进行计算,而Q使用上一个 Decoder block 的输出计算。
Masked MHA: 与MHA类似,只是多了一个掩码机制,Transformer中有两个掩码机制:
(1)padding mask:输入序列对齐,对短序列尾部补0,长序列截断,具体的做法就是对阶段的位置加上一个非常大的负数,这样经过softmax,这些位置的概率就会接近0。attention机制就不会把注意力放在这些位置上;(Encoder也需要这一步)
(2)sequence mask:阻止decoder看见未来的信息,将time_step = t之后的信息隐藏起来,通过一个上三角矩阵作用在序列上即可。
需要注意的是,Decoder的第二个MHA就不是self-attention了,其query来自第一个Masked MHA,k,v来自Encoder的输出;那么第一个Masked MHA的输入Output Embedding是什么呢?他实际上就是一个根据Decoder输出逐渐右移的输出序列,以输入标识符
1.3 代码
点击查看代码
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class LayerNorm(nn.Module):
"""Layer Normalization"""
def __init__(self, dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.bias = nn.Parameter(torch.zeros(dim))
self.eps = eps
def forward(self, x):
mean = x.mean(-1, keepdim=True)
std = x.std(-1, keepdim=True)
return self.weight * (x - mean) / (std + self.eps) + self.bias
class SelfAttention(nn.Module):
"""
自注意力模块(支持全局和因果两种模式)
"""
def __init__(self, dim, num_heads=8, dropout=0.0, causal=False):
super().__init__()
assert dim % num_heads == 0, "dim must be divisible by num_heads"
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.causal = causal # True=Stage2, False=Stage1
# Q, K, V 投影
self.qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
self.dropout = nn.Dropout(dropout)
# 因果掩码(Stage 2 用)
if causal:
self.register_buffer(
'causal_mask',
torch.tril(torch.ones(1024, 1024)).view(1, 1, 1024, 1024)
)
def forward(self, x, mask=None):
"""
Args:
x: [B, N, D] 输入序列
mask: [B, N] 可选的 padding mask
Returns:
out: [B, N, D] 输出序列
"""
B, N, D = x.shape
# 1. 生成 Q, K, V [B, N, 3*D] -> [3, B, Heads, N, HeadDim]
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
# 2. 计算注意力分数 [B, Heads, N, N]
scale = 1.0 / math.sqrt(self.head_dim)
attn = (q @ k.transpose(-2, -1)) * scale
# 3. 应用因果掩码(Stage 2)
if self.causal:
causal_mask = self.causal_mask[:, :, :N, :N]
attn = attn.masked_fill(causal_mask == 0, float('-inf'))
# 4. 应用 padding mask(可选)
if mask is not None:
mask = mask.view(B, 1, 1, N)
attn = attn.masked_fill(mask == 0, float('-inf'))
# 5. Softmax + Dropout
attn = attn.softmax(dim=-1)
attn = self.dropout(attn)
# 6. 加权求和 [B, Heads, N, HeadDim]
out = (attn @ v).transpose(1, 2).reshape(B, N, D)
# 7. 输出投影
out = self.proj(out)
return out
class MLP(nn.Module):
"""
前馈神经网络 (Feed Forward)
"""
def __init__(self, dim, hidden_dim=None, dropout=0.0):
super().__init__()
hidden_dim = hidden_dim or dim * 4
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class TransformerBlock(nn.Module):
"""
单个 Transformer Block (Pre-Norm 架构)
"""
def __init__(self, dim, num_heads=8, mlp_ratio=4.0, dropout=0.0, causal=False):
super().__init__()
self.norm1 = LayerNorm(dim)
self.attn = SelfAttention(dim, num_heads, dropout, causal=causal)
self.norm2 = LayerNorm(dim)
self.mlp = MLP(dim, int(dim * mlp_ratio), dropout)
def forward(self, x, mask=None):
# Pre-Norm + 残差连接
x = x + self.attn(self.norm1(x), mask=mask)
x = x + self.mlp(self.norm2(x))
return x
class TransformerEncoder(nn.Module):
"""
完整的 Transformer Encoder(Stage 1 用,全局注意力)
"""
def __init__(self, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.blocks = nn.ModuleList([
TransformerBlock(dim, num_heads, mlp_ratio, dropout, causal=False)
for _ in range(depth)
])
self.norm = LayerNorm(dim)
def forward(self, x, mask=None):
"""
Args:
x: [B, N, D] 输入序列
mask: [B, N] 可选的 padding mask
Returns:
x: [B, N, D] 输出序列
"""
for blk in self.blocks:
x = blk(x, mask=mask)
return self.norm(x)
class TransformerDecoder(nn.Module):
"""
完整的 Transformer Decoder(Stage 2 用,因果注意力)
"""
def __init__(self, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.blocks = nn.ModuleList([
TransformerBlock(dim, num_heads, mlp_ratio, dropout, causal=True)
for _ in range(depth)
])
self.norm = LayerNorm(dim)
def forward(self, x):
"""
Args:
x: [B, N, D] 输入序列(已包含位置编码)
Returns:
x: [B, N, D] 输出序列
"""
for blk in self.blocks:
x = blk(x) # 因果掩码在 Attention 内部处理
return self.norm(x)
1.4 复杂度与显存占用计算
自注意力模块主要考虑线性变换、注意力计算、加权求和三部分:
线性变换:需要计算QKV三个线性变换,每个变换的复杂度为\(O(n \cdot d^2_{model})\),其中\(d_{model}\)是隐藏层维度,n是序列长度。
注意力计算:计算\(QK^T\)的复杂度为\(O(n^2 \cdot h \cdot d_k)\),其中h是头数,\(d_k\)为每个头的维度。
加权求和:计算加权值并进行最终的线性变换,复杂度为\(O(2n^2 \cdot d_{model} + 3n \cdot d^2_{model})\)
因此,总计算复杂度为:

前馈网络主要需要计算前置与后置的线性变换,激活函数:类似的,有\(2 * O(n \cdot d_{model} \cdot d_{ff} + O(1)\)。
一般来说,自注意力因为有二次项,计算成本显著高于前馈网络,若\(d_{ff}\)很大时,如\(d_{ff} = 4 \cdot d_{model}\),其计算量可能接近自注意力,但依然低于二次项。
显存占用主要看中间结果的存储消耗,在自注意力模块中,主要包含Q、K、V矩阵,注意力矩阵\(QK^T\)和加权求和中间结果,其中Q、K、V矩阵每个维度为\(O(n \cdot d_{model}\),注意力矩阵\(QK^T\)维度为\(n \times n\),显存为\(n^2\),中间值忽略不计,总显存占用为\(O(n^2 + n \cdot d_{model}\)
前馈网络同理,中间层激活为\(O(n \cdot d_{ff}\),输入和输出为\(O(2 \cdot n \cdot d_{model}\)。总显存占用就是两者相加,主要看{d_{ff}}与{n}的数量级关系
可以看到,显存占用主要来自注意力矩阵,序列长度越长,显存需求增长的就越快。{d_{ff}}通常是我们自行设置的,如果设置的过大,可能会超过自注意力,这是需要注意的一点。一般来说,GPT3这类长文本生成注意力矩阵的问题是主要挑战,对于BERT来说,前馈网络中模型维度的设置存在隐性的显存问题。实际应用中,需根据模型参数(如\(d_{ff}\)、n)和硬件限制选择优化策略(如分块计算注意力、减少\(d_{ff}\)或使用稀疏注意力)。
2. 大模型的结构设计
以llama3.1-405B为例,llama系列本质上仍然是对Decoder-only Transformer架构的组件进行改进的一个基底大语言模型,具体的结构框架示意图引用了知乎张俊林的一篇解读:

Llama 模型的架构演变主要经历了以下三个阶段:
Llama 1:基于原始 Transformer 架构,引入了预归一化、RMSNorm、SwiGLU 激活函数和旋转式位置编码等改进,提升了模型的训练稳定性和性能。
Llama 2:在 Llama 1 的基础上,将上下文长度扩展至 4096,并引入了分组查询注意力 (GQA) 机制,有效降低了推理过程中的内存需求,提升了推理速度。
Llama 3:进一步将 GQA 应用于小型模型,并采用更高效的分词器 TikToken,扩大了词汇表的数量,同时将上下文长度翻倍,并大幅增加了训练数据量。
llama 1
llama 1借鉴了GPT3架构中提高训练稳定性的方法,也采用了对每个Transformer子层的输入进行归一化的策略,而不仅仅是对输出进行归一化处理。(即Embedding后先Add&Norm,然后再正常MHA、Add&Norm,FFN,Add&Norm)
此外,还采用RMSNorm来替换传统的LayerNorm,这一改变在保持训练稳定性和提升模型收敛速度的同时,大幅提高了计算效率。除此之外,在激活函数的选择上,llama 1采用了SwiGLU函数来替代传统的ReLU函数,这一改变旨在提升模型的性能。
ReLU会将所有复数输入直接归零,而正数输入则保持不变。
相比之下,SwiGLU函数含有一个可学习的参数\(\beta\),能够调节函数的插值程度。随着\(\beta\)值的增大,SwiGLU的行为将逐渐接近ReLU。(然而,SwiGLU的成功原因实际并没有给出解释,我们一般认为,1、Swish对于负值的响应相对较小克服了 ReLU 某些神经元上输出始终为零的缺点;2、GLU 的门控特性,这意味着它可以根据输入的情况决定哪些信息应该通过、哪些信息应该被过滤。这种机制可以使网络更有效地学习到有用的表示,有助于提高模型的泛化能力。在大语言模型中,这对于处理长序列、长距离依赖的文本特别有用;3、SwiGLU 中的参数 W1,W2,W3,b1,b2,b3W1,W2,W3,b1,b2,b3 可以通过训练学习,使得模型可以根据不同任务和数据集动态调整这些参数,增强了模型的灵活性和适应性;4、计算效率相比某些较复杂的激活函数(如 GELU)更高,同时仍能保持较好的性能。这对于大规模语言模型的训练和推理是很重要的考量因素。)
点击查看代码
class SwiGLU(nn.Module):
def __init__(self, w1, w2, w3) -> None:
super().__init__()
self.w1 = w1
self.w2 = w2
self.w3 = w3
def forward(self, x):
x1 = F.linear(x, self.w1.weight)
x2 = F.linear(x, self.w2.weight)
hidden = F.silu(x1) * x2
return F.linear(hidden, self.w3.weight)
旋转位置编码也是一个重要改进点,通过对序列每个位置转换成词嵌入的旋转变量来模拟单次间的相对位置关系,采用这种方式,即便在原句中增加更多词汇,单词之间的相对距离也能得到保持。
llama 2
llama 2主要基于llama 1进行修改,将处理上下文的长度扩展至4096(原为2048),对于34B和70B的大型模型,llama 2使用Grouped-Query Attention取代了原本的MHA,即g组Q,每组Q共享1个KV。
除此以外,llama更多是对于预训练数据集增大了40%,并对对齐技术,特别是SFT与RLHF部分进行了着重微调,保持LLM与对话式指令的一致。具体的介绍将在第三章详细节能型。
llama 3
Llama 3 将处理上下文的长度从 4096 扩展至 8192,并将 GQA 使用到了较小规模的模型(8B)。同时,研究者们还将分词工具从 Sentence Piece 更换为 OpenAI 模型所采用的 TikToken。因为新的词汇表容量增加到了 128k 个 tokens,较之前的 32k 有了大幅提升,这一变更显著提升了模型的性能。
这两种分词工具的主要差异在于,在输入的 tokens 已经存在于词汇表中时,TikToken 会跳过字节对编码(BPE)的合并规则。例如,如果“generating”这个词已经在词汇表中了,那么它将作为一个完整的 token 返回,而不是将其拆分为“generating”和“ing”这两个最小单元的 tokens 。
llama 3.1 实现了上下文长度(128K tokens)的显著提升,并新增了对 8 种语言的支持。此次发布版本的一个重要亮点是更大的 Llama 3.1 405B 模型。在此之前,开放式的 LLMs(大语言模型)通常模型规模都低于 100 B。其中讨论提出:可以利用强大的教师模型来创建性能更佳的小型模型。
llama 3.2 11B 和 90B 支持图像推理用例,例如文档级理解(包括图表和图形)、图像字幕以及视觉基础任务(例如基于自然语言描述在图像中精确定位对象)。轻量级1B和3B模型具有强大的多语言文本生成和工具调用功能。这些模型使开发人员能够构建个性化的设备代理应用程序,具有很强的隐私性,数据永远不会离开设备。 例如,这样的应用程序可以帮助汇总最近收到的 10 条消息,提取操作项,并利用工具调用直接发送日历邀请以进行后续会议。
视觉模型创新:为了添加图像输入支持,训练了一组适配器权重,将预训练的图像编码器集成到预训练的语言模型中。适配器由一系列交叉注意层组成,这些层将图像编码器表示输入到语言模型中。 在文本-图像对上训练了适配器,以使图像表示与语言表示对齐。在适配器训练期间,还更新了图像编码器的参数,但有意不更新语言模型参数。通过这样做,保持了所有纯文本功能不变,为开发人员提供了 Llama 3.1 模型的直接替代品。
Llama 3.2训练流程由多个阶段组成,从预训练的 Llama 3.1 文本模型开始:
添加图像适配器和编码器
在大规模噪声(图像、文本)对数据上进行预训练。
在中等规模的高质量领域内和知识增强的(图像、文本)对数据上进行训练。
在后期训练中,使用与文本模型类似的方法,在监督微调、拒绝采样和直接偏好优化方面进行多轮对齐。
利用 Llama 3.1 模型生成合成数据,在域内图像的基础上过滤和扩充问题和答案,并使用奖励模型对所有候选答案进行排名,以提供高质量的微调数据。还添加了安全缓解数据,以生成具有高安全水平的模型,同时保留模型的有用性
最终结果是一组可以同时接收图像和文本提示并深入理解和推理两者组合的模型。这是 Llama 模型向拥有更丰富代理能力迈出的又一步。
同时,1B和3B模型使用了教师模型进行了修建和提炼,使其成为首批能够高效适应设备的高性能轻量级 Llama 模型。具体介绍将放在第三章进行介绍。

以上文本参考:> https://zhuanlan.zhihu.com/p/8828758250 > https://zhuanlan.zhihu.com/p/760251326

最后,在稍微介绍一下大模型能力的评估方法。
评估包含三个量化指标:第一个是评估推理速度,第二个是确定答案长度,第三个是评估准确性。对于准确性的评估,可以使用RAQ。RAQ 通过一个独立的 LLM 对 Llama 2 和 Llama 3 的答案进行排序,排序的依据是它们与真实答案的接近程度。当然,最近ICML2026里的一篇文章指出,用大模型judge大模型会出现judge不一致的问题,一个解决办法就是最后不要直接输出softmax后的判断,而是直接进行离散概率的输出,然后基于每个位置的判断再统一进行judge即可。
不同语言模型性能可以用 Dunn 事后检验(Dunn post-hoc test)结果进行评估。
3. 模型微调
为什么需要模型微调,一个最大的原因就是,直接重新训练一个大模型的代价很大。
一个6B的大模型如果要进行全量微调,需要多少显存呢?
全量微调,顾名思义,就是加载预训练好的模型权重,然后在我们自己的特定任务数据上,更新模型的所有参数。听起来很简单,但魔鬼藏在细节里。在训练过程中,显存主要被以下几个部分占用:
模型参数(Model Parameters)
梯度(Gradients)
优化器状态(Optimizer States)
激活值(Activations)
临时缓冲区和输入数据(Temporary Buffers & Input Data)
对于FP32,模型参数本身的6B,需要6B * 4 bytes = 24 GB,FP16可以减少一半,是目前训练和微调 LLM 最常用的精度。INT8 可以再减少一半。
反向传播过程中,模型的每一个可训练参数都需要计算梯度,因此二者参数量相同,精度也一致,故显存相同24/12/6 GB。
优化器状态:Adam/AdamW 通常需要存储一些中间状态,以便更有效地更新参数。对于Adam来说,需要为每个模型参数存储两个状态:动量、方差。这两个状态通常与参数/梯度精度相同,特别地,为了训练稳定性,即使模型参数和梯度使用FP16,优化器状态通常会保持在FP32.这是一种常见的混合精度训练策略。那么优化器状态需要24+24=48GB,即便你牺牲一些性能强行做FP16,也要24GB。
激活值是比较难以估算的显存占用部分,取决于批量大小 (Batch Size - B): 每个 GPU 处理的样本数。
序列长度 (Sequence Length - S): 输入序列的最大 Token 数。
模型隐藏层维度 (Hidden Dimension - H)
模型层数 (Number of Layers - L)
模型架构细节 (Attention机制等)
影响: 复杂且依赖具体实现。
解释: 不同的架构或同一架构的不同实现方式会影响需要存储哪些具体的中间值以及它们的精确大小。
自注意力 (Self-Attention): 需要存储 Q, K, V 矩阵,注意力输出等。如前所述,注意力得分矩阵可能是个巨大的临时存储。
前馈网络 (Feed-Forward Network - FFN): 中间扩展层的激活值是主要的显存消耗者。
归一化层 (Normalization Layers): 如 LayerNorm,也需要存储其计算过程中的中间统计量(均值、方差)或输入,以便反向传播。
残差连接 (Residual Connections): 需要存储添加残差之前的那个分支的输出。
其他复杂结构: 如 MoE (Mixture of Experts) 等会引入更复杂的激活值存储模式。
是否使用梯度检查点 (Gradient Checkpointing / Activation Recomputation)
激活值是在前向传播过程中计算出来的中间结果,需要暂存下来供反向传播时计算梯度。其大小约等于 B * S * H * L * c * (激活值存储精度对应的字节数)。
对于一个 6B 参数的 Transformer 模型(类似 Llama 架构),H 可能在 4096 左右,L 在 32 层左右。
假设使用 BF16/FP16 (2 bytes) 存储激活值。
假设一个适中的 Batch Size (B=4) 和 Sequence Length (S=2048)。
C 是一个常数因子,代表平均每层每个 Token 每个隐藏单元需要存储的激活值数量。这个因子非常依赖架构。对于 Transformer,它不仅包括层输出 (B, S, H),还包括 FFN 的中间扩展层 (B, S, intermediate_size),以及 Q, K, V 等。C 可能在 10 到 24 甚至更高,取决于具体实现和计算方式。
一个非常粗略的经验法则是,激活值内存可能与模型参数、梯度、优化器状态的总和处于相似的数量级,甚至更大。
无梯度检查点: 激活值内存可能非常巨大。对于 B=4, S=2048, H=4096, L=32, BF16 的配置,激活值可能轻松达到 几十 GB甚至上百 GB。例如,粗略估算可能在 40 GB - 100 GB 之间,具体取决于实现细节和模型结构。
使用梯度检查点: 这项技术通过在反向传播时重新计算部分激活值,而不是全部存储,来显著减少显存占用,但会增加计算时间(约 20-30%)。使用梯度检查点后,激活值内存占用可以大幅降低,可能降至 5 GB - 20 GB 的范围,具体取决于检查点的设置策略。
最后的临时缓冲区与输入数据包括输入数据本身(Token IDs, Attention Masks 等)转换成 Embedding 后的显存占用、CUDA Kernel 运行时可能需要的临时显存、各种计算过程中产生的中间变量。这部分通常相比前面几项要小,但也不能完全忽略,尤其是在显存非常紧张的情况下。
因此,总的显存量,以FP16为例,需要至少72+40100(520)GB。
那么,大模型最重要的一个核心思想:PEFT参数高效微调,就出现了。PEFT的核心思想是:冻结(Freeze) 预训练大模型的大部分(甚至几乎全部)参数,只微调其中少量或额外增加的一小部分参数。
这样做的好处是显而易见的:
大幅降低显存需求:
梯度: 只需为少量可训练参数计算和存储梯度。
优化器状态: 只需为少量可训练参数存储优化器状态。
模型参数本身(占大头)虽然加载了,但不产生梯度和优化器状态,极大地节省了显存。
激活值部分依然存在,但由于可训练参数减少,整体显存压力骤减。
减少计算量,加快训练速度: 反向传播和参数更新的计算量大大减少。
降低存储成本: 每个任务只需要存储少量修改的参数,而不是整个模型的副本。
可能获得更好的性能: 有研究表明,PEFT 方法有时能比全量微调在新任务上表现更好,尤其是在低数据场景下,可以有效防止灾难性遗忘(Catastrophic Forgetting),并可能具有更好的泛化能力。
根据其实现方式,PEFT 技术大致可以分为三类:
增加额外参数 (Adding Extra Parameters, A): 在原有模型结构中插入或附加少量新的、可训练的模块或参数,冻结原始参数。
类适配器 (Adapter-like) 方法: 如 Adapter Tuning,在 Transformer 层之间插入小的瓶颈结构(Adapter 模块)。
软提示 (Soft Prompts): 如 Prefix Tuning, Prompt Tuning, P-Tuning 等。它们不是调整模型权重,而是在输入或中间层添加可训练的连续向量(Virtual Tokens 或 Prompts),引导模型在下游任务上表现。
选取一部分参数更新 (Selecting Parameters, S): 只选择模型中的一部分已有参数进行微调。
例如 BitFit,它只微调模型中的 Bias(偏置)参数,或者只微调某几层。
引入重参数化 (Reparameterization, R): 利用一些参数化的方式来表示参数的更新,从而减少需要直接优化的参数量。
LoRA 也可以看作是重参数化的一种形式。

以上提到的 BitFit、Prefix Tuning、Prompt Tuning、P-Tuning、Adapter Tuning、LoRA 等都是当前非常热门且有效的 PEFT 方法。它们各自有不同的设计哲学和适用场景。
1. 添加附加参数(Additive Methods)
这类方法的核心思想是冻结原始模型的绝大部分参数,然后添加一些新的、参数量较小的模块或参数,在训练时只更新这些新添加的部分。
Adapter Tuning (适配器微调):
原理:这是比较早期的PEFT方法。它在Transformer的每个层(或部分层)中插入两个小的、瓶颈结构(bottleneck)的前馈神经网络模块(即Adapter模块)。通常是先降维,经过一个非线性激活函数,再升维恢复到原来的维度。
形象比喻:就像在预训练模型的“信息高速公路”的每个收费站旁边,加了一个小小的“咨询台”。信息流过时,会先去咨询台“问问路”(经过Adapter模块处理),然后再继续前进。训练时,我们只训练这些咨询台的工作人员。
优点:模型结构清晰,易于实现。
缺点:可能会增加一点点推理时的延迟,因为信息流需要额外经过这些小模块。
Prefix-Tuning / Prompt-Tuning / P-Tuning系列:
原理:这类方法受到人类如何通过提示(Prompt)引导模型生成内容的启发。它们不是修改模型内部结构,而是在输入层或者模型的注意力层添加一些可学习的连续向量(称为Virtual Tokens或Prefixes)。模型在处理任务时,会把这些可学习的向量和正常的输入一起处理,从而引导模型的行为。
Prompt-Tuning:只在输入嵌入层添加可学习的Tokens。
Prefix-Tuning:在每一层的Key和Value向量前添加可学习的Prefixes。
P-Tuning:使用一个小的LSTM或MLP来生成这些可学习的Tokens,使其更具上下文关联性。P-Tuning v2则进一步优化,应用到更深层。
形象比喻:想象你在给大模型下达指令。Prefix/Prompt Tuning就像是在你的指令前面加上一句“魔法咒语”(可学习的向量),这句咒语能引导模型更好地理解和执行你的特定任务,而我们只需要学习这句“咒语”怎么念。
优点:需要更新的参数量极少(有时只有几十万甚至几万),效果也不错。
缺点:训练有时不如Adapter或LoRA稳定,性能上限可能稍低(尤其早期版本)。
2. 选择性微调(Selective Methods)
这类方法选择性地只微调模型中的一部分已有参数,比如只微调Bias项(偏置参数),或者只微调顶部的几层。
原理:基于假设——模型的部分参数对适应新任务更重要。例如,Bias参数通常被认为与模型的风格或领域适应性关联更强。
现状:虽然思路直接,但在实践中,其效果和普适性往往不如Additive或Reparameterization方法,目前相对不是主流的PEFT方向,但其思想仍有借鉴意义。
BitFit:仅调整模型中的偏置项(Bias Terms),其他参数全部冻结。实验表明,BERT-large仅需调整0.1%的参数即可在GLUE任务中达到95%的全量微调性能14。
DiffPruning:通过稀疏化方法动态选择需微调的参数子集,减少计算冗余。
3. 基于重参数化(Reparameterization-based Methods)
这类方法通过引入低秩分解等技术来“间接”地修改模型的参数。
LoRA (Low-Rank Adaptation):
原理:这是目前最流行、应用最广泛的PEFT方法之一。LoRA的核心假设是:模型在适应新任务时,其参数矩阵的变化(ΔW)是低秩的。也就是说,这个变化可以用两个更小的矩阵(A和B)的乘积来近似表示 (ΔW ≈ BA)。在微调时,原始权重W保持冻结,我们只训练这两个小的低秩矩阵A和B。推理时,可以将学习到的BA加回到原始权重W上(W’ = W + BA),这样不会引入额外的推理延迟。
形象比喻:想象原始模型的参数矩阵W是一张巨大的、复杂的高清地图。全量微调是重新绘制整张地图。LoRA则是认为,针对新任务的“路线更新”,只需要在原地图上贴几张小的“修正贴纸”(BA),而我们只需要学习制作这些贴纸(训练A和B)。贴纸很小,制作起来很快。用的时候,可以直接把贴纸固定在地图上。
优点:效果通常很好,训练效率高,推理时可以合并参数(无额外延迟),实现简单。
缺点:秩(Rank)的选择是一个超参数,需要调整。
QLoRA (Quantized LoRA):
原理:LoRA的进一步优化。它在LoRA的基础上,结合了模型量化(Quantization)技术。在训练时,将冻结的预训练模型参数量化到更低的数据类型(比如4-bit),大大减少了训练时的显存占用。同时,LoRA部分仍然用较高精度(如16-bit)进行训练,保证了微调的精度。
优点:极大降低了显存门槛,使得在消费级GPU上微调非常大的模型成为可能。
缺点:实现相对复杂一些,需要处理量化和反量化的细节。
PEFT的发展路径与趋势
回顾PEFT的发展,我们可以看到一条清晰的脉络:
早期探索(~2019-2020):以Adapter为代表,验证了“只动一小部分”的可行性。
Prompt范式兴起(~2021):Prefix-Tuning、Prompt-Tuning等方法展示了通过“外部引导”进行高效微调的可能性,参数量进一步降低。
LoRA异军突起(~2021-至今):LoRA凭借其出色的效果、效率和易用性,迅速成为主流,并衍生出QLoRA等众多改进版本,极大地推动了PEFT技术的普及。
融合与未来:目前的研究趋势包括:
组合方法:尝试结合不同PEFT方法的优点。
自动化选择:研究如何自动确定哪些参数最值得微调。
更极致的效率:探索在更低资源下实现高性能微调的方法。
理论理解:深入探究PEFT为何有效,以及不同方法背后的数学原理。
如何准备面试中的PEFT问题?
如果你在面试中被问到PEFT,面试官可能想了解:
你知道为什么要用PEFT吗? (回答:解决全量微调的资源消耗、存储、部署、遗忘问题)
你知道哪些主流的PEFT方法? (至少能说出Adapter, Prompt-Tuning/Prefix-Tuning, LoRA)
能简单解释一下LoRA的原理吗? (回答:冻结原模型,用低秩矩阵分解近似参数变化,只训练低秩矩阵A和B,推理时可合并)
LoRA相比Adapter有什么优势? (通常效果更好,推理时可合并无延迟)
你知道QLoRA吗?它解决了什么问题? (回答:结合量化,极大降低显存占用)
PEFT方法之间有哪些主要的trade-off? (性能 vs. 参数量 vs. 训练/推理效率 vs. 实现复杂度)
除此以外,我们再介绍一下llama中引入的SFT和RLHF。
SFT
SFT是在预训练模型基础上,使用高质量标注的instruction-response对进行有监督微调,让模型学会遵循指令、完成特定任务
点击查看代码
# LLaMA 2/3采用的Chat Template格式
<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a helpful assistant<|eot_id|>
<|start_header_id|>user<|end_header_id|>
请解释Transformer的注意力机制<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
Transformer通过Query-Key-Value计算注意力权重...<|eot_id|>
这一阶段的微调使用QLoRA(4-bit量化+LoRA)可在单卡24G显存微调8B模型,开启packing将多条短样本拼接,提升训练效率30%+,采用gradient checkpointing节省显存,换取少量计算时间。
虽然SFT能让模型学会执行任务,但它本质是模仿训练数据分布。如果数据存在偏见,模型会放大偏见;如果遇到训练分布外的复杂推理,模型容易'一本正经胡说八道'。这就是为什么需要RLHF做二次对齐。
RLHF
┌─────────────────────────────────────────┐
│ Stage 1: Reward Model Training │
│ • 收集人类对同一prompt多个response的排序│
│ • 训练RM学习"人类偏好"的打分函数 │
│ • 关键:标注一致性>80%,避免噪声信号 │
└─────────────────────────────────────────┘
▼
┌─────────────────────────────────────────┐
│ Stage 2: PPO Policy Optimization │
│ • 用SFT模型作为initial policy │
│ • RM给出reward,PPO更新policy参数 │
│ • 加入KL散度约束,防止偏离SFT太远 │
└─────────────────────────────────────────┘
▼
┌─────────────────────────────────────────┐
│ Stage 3: Iterative Refinement │
│ • Llama 2采用V1→V5多轮迭代[[21]][[22]] │
│ • 每轮收集新反馈,持续优化对齐效果 │
└─────────────────────────────────────────┘
是否需要模型具备专业领域知识?
├─ 是 → 优先SFT(用领域instruction数据微调)
└─ 否 → 进入下一判断
是否涉及主观偏好/价值观对齐?
├─ 是 → 必须RLHF(如安全、无害、有帮助)
└─ 否 → SFT可能足够
资源是否受限?
├─ 是 → QLoRA-SFT + DPO替代完整RLHF[[3]][[9]]
└─ 否 → 完整SFT+RLHF pipeline
数据标注成本是否可控?
├─ 标注response对 → SFT
├─ 标注response排序 → RLHF的RM训练
└─ 无标注数据 → 考虑无监督/自监督方法
4. 分布式训练
Data Parallelism
数据并行(DP)是最直观、最容易理解的一种并行方式。想象一下,你要训练一个模型,但是数据量太大了,一个人处理不过来。怎么办?找一群人(GPU),每个人都拿到一份完整的操作手册(模型),然后把任务(数据集)分成几份,每人领一份去干活(训练)。具体流程:
分发: 把一个大的Batch数据切分成N个小的mini-batch,每个GPU分到一个mini-batch。
复制: 每个GPU上都保留一份完整的模型副本。
计算: 每个GPU独立地根据自己的mini-batch数据进行前向传播计算损失,然后进行反向传播计算梯度。
同步: 这是关键!在反向传播之后,所有GPU计算出的梯度需要进行一次聚合(Aggregation),通常是做一次All-Reduce操作(所有GPU把自己的梯度发给其他所有GPU,并计算平均值),保证所有GPU上的模型副本在参数更新后仍然保持一致。
更新: 每个GPU使用聚合后的梯度来更新自己的模型参数。
优点:
实现简单,容易理解和上手。PyTorch的DistributedDataParallel (DDP) 就是典型的实现,几行代码就能搞定。
通常能获得不错的加速比,尤其是在网络带宽足够的情况下。
适用于数据量大,但模型可以放入单张 GPU 的情况。
缺点:
内存冗余: 每个GPU都要存一份完整的模型参数、梯度和优化器状态。当模型非常大时,即使一个副本也可能塞不进单张卡的显存,这时候DP就无能为力了。
面试小贴士: DDP是面试常考点。可以提一下它相比于更早的DataParallel(DP) 的优势。
DP的工作方法是:
-
基于单进程多线程。它使用 Python 的 GIL (Global Interpreter Lock),这意味着在同一时刻只有一个线程能执行 Python 字节码。在涉及大量 Python 逻辑或数据加载时,GIL 会成为瓶颈。
-
通信后端使用 Python 线程进行梯度聚合,通信效率较低。
-
GPU 0 承担额外的工作:它不仅要计算自己的前向/反向传播,还要负责聚合所有其他卡的梯度、更新模型参数,然后将新参数广播给其他卡。这导致 GPU 0 的显存占用比其他卡高,计算负载也更重,容易成为短板(Straggler),拖慢整体速度。仅支持单机多卡。它无法跨越多台服务器进行训练。
因此相比之下,DDP使用多进程,每个 GPU 对应一个独立的进程,完全绕过 GIL 限制。每个进程独立加载数据、计算梯度,效率更高。避免了Python GIL(全局解释器锁)的瓶颈,并且通信效率更高(直接进行All-Reduce,而不是先把梯度汇总到主GPU再分发)。
Model Parallelism
模型并行,顾名思义,就是把一个大模型拆分成多个部分,分别放到不同的GPU上。这样每个GPU只需要承担模型的一部分,显存压力就小多了。
模型并行主要有两种玩法:张量并行(Tensor Parallelism, TP) 和 流水线并行(Pipeline Parallelism, PP)。你可以粗略地理解为:张量并行是在层内拆,流水线并行是在层间拆。
张量并行主要针对模型中的单个大运算(比如巨大的nn.Linear层或Attention层里的矩阵乘法)进行拆分。它把一个大的张量(Tensor)沿着某个维度切成N块,每个GPU只持有和计算其中的1/N块。

Transformer中的应用: 在Transformer的MLP(多层感知机)部分,通常是先一个nn.Linear(记为A),然后接一个激活函数,再接一个nn.Linear(记为B)。Megatron-LM 提出了一种巧妙的1D张量并行方法:
第一个Linear层A按列切分。计算时,输入X需要通过All-Gather在TP组内同步,然后每个GPU计算 X * Ai,得到 Yi。
第二个Linear层B按行切分。输入是Y = [Y0, Y1, ..., Yn-1],每个GPU持有Yi和Bi。计算 Zi = Yi * Bi。最后,需要对所有GPU的输出 Zi 进行一次All-Reduce,得到最终结果 Z = sum(Zi)。
通过这种方式,两次大的矩阵乘法都被有效地并行化了,并且巧妙地安排了通信操作(一次All-Gather,一次All-Reduce)。
优点:
可以有效降低单张卡上的峰值显存(尤其是激活值的显存)。
允许训练单个层就非常巨大的模型。
缺点:
通信开销大: 需要在每次前向和反向传播过程中进行额外的通信(如All-Gather, All-Reduce)。
实现复杂: 需要深入模型内部,修改算子的实现。
通常局限于节点内: 因为通信量大,对网络带宽要求极高,一般只在节点内部(比如通过NVLink高速互联的GPU)使用。
p.s.张量并行和数据并行的区别。核心在于切分的对象不同(模型 vs 数据)以及通信模式不同(TP通信更频繁,通常在算子内部)。可以结合Transformer的MLP或Attention解释TP的具体切分和通信过程。
流水线并行,Pipeline Parallelism
把模型按层分组,不同组交给不同GPU,串行处理。
流水线并行是把模型的不同层分配到不同的GPU上。比如一个有40层的模型,用4张卡做流水线并行,可以把0-9层放GPU 0,10-19层放GPU 1,以此类推。
朴素流水线的问题(Naive Pipeline Parallelism):
最简单的想法是:一个batch的数据先在GPU 0上完成前10层的前向计算,然后把输出(激活值)传给GPU 1;GPU 1计算10-19层,再传给GPU 2…直到最后一个GPU完成计算。反向传播类似地倒着来一遍。
这种方式的问题在于“流水线气泡(Bubble)”:在任何一个时间点,只有一个GPU在干活,其他GPU都在等待。这就像一个效率低下的工厂流水线,工人大部分时间在发呆。
改进:微批次(Micro-batching)
为了减少气泡,提高设备利用率,现代流水线并行(如GPipe, PipeDream)引入了微批次(Micro-batch)的概念。将一个大的训练批次(mini-batch)进一步切分成更小的微批次。然后,让这些微批次像流水一样流过各个GPU阶段(Stage)。
当前一个GPU处理完一个微批次后,立刻将其输出传给下一个GPU,然后自己开始处理下一个微批次。这样,在理想情况下,多个GPU可以同时处理不同的微批次,大大减少了空闲时间。
GPipe vs PipeDream (1F1B):
GPipe: 先把所有微批次的前向传播都做完,然后再做所有微批次的后向传播。这种方式需要缓存每个微批次在前向传播时的激活值,内存开销较大。
PipeDream (及其变种,如1F1B - One Forward, One Backward): 采用更优化的调度策略,例如“1F1B”。一个GPU完成一个微批次的前向计算后,一旦后续阶段也完成了该微批次的前向,并且计算出了梯度,这个梯度就会被立刻传回给当前GPU,使其可以尽快开始计算该微批次的后向传播。这样可以显著减少需要缓存的激活值数量,降低内存压力。
优点:
显著降低单卡内存占用(只存部分层和少量激活值)。
可以训练非常深的模型。
相比TP,对节点间带宽要求稍低。
缺点:
存在流水线气泡,理论加速比不如数据并行。气泡大小和微批次数有关,需要调优。
负载均衡问题:需要尽量让每个阶段的计算量和参数量均衡,否则慢的阶段会成为瓶颈。
实现相对复杂,需要对模型进行切割,处理跨设备的数据传输。
实现: GPipe, PipeDream (及其变种), Megatron-LM, DeepSpeed Pipeline Parallelism。
优化器状态并行/ZeRO:
训练时的显存占用不仅有模型参数,还有梯度和优化器状态(尤其是Adam这种带一阶、二阶动量的优化器,状态量可能是模型参数的两倍!)。数据并行虽然加速了计算,但这些状态在每个GPU上都存了一份,造成了巨大的内存冗余。
ZeRO(Zero Redundancy Optimizer) 就是来解决这个问题的。
核心思想: 既然冗余,那就分片(Shard),每人只存一部分。
ZeRO本质上是一种增强的数据并行,它的目标是消除数据并行中的内存冗余。它不是复制完整的模型状态(参数、梯度、优化器状态),而是将这些状态分割成N份(N是参与ZeRO的GPU数量),每个GPU只负责存储和更新其中的1/N份。

ZeRO的三个阶段:
ZeRO-Stage 1 (ZeRO-1):优化器状态分片 (Optimizer States Sharding)
每个GPU只保存和更新它负责的那部分参数对应的优化器状态。
模型参数和梯度在每个GPU上仍然是完整的副本。
在优化器更新步骤(optimizer.step())时,需要一次通信(Reduce-Scatter)将梯度聚合到对应的GPU上,然后每个GPU更新自己负责的那部分参数对应的优化器状态和参数。之后再通过一次通信(All-Gather)将更新后的完整参数同步给所有GPU。
效果: 显著减少优化器状态的内存占用(约减少 (N-1)/N)。
ZeRO-Stage 2 (ZeRO-2):优化器状态 + 梯度分片 (Optimizer States & Gradients Sharding)
在ZeRO-1的基础上,梯度也进行分片。
在反向传播计算梯度时,每个GPU只保留自己负责的那部分参数对应的梯度,其他参数的梯度计算完后就丢弃。通过一次Reduce-Scatter操作,将梯度直接聚合到目标GPU上。
模型参数在每个GPU上仍然是完整的副本(在前向/反向计算时需要)。
效果: 进一步减少梯度占用的内存(约减少 (N-1)/N)。优化器更新时的通信模式与ZeRO-1类似。
ZeRO-Stage 3 (ZeRO-3):优化器状态 + 梯度 + 模型参数分片 (Optimizer States & Gradients & Parameters Sharding)
终极形态!把模型参数也分片了。每个GPU平时只持有自己负责的那部分模型参数。
关键操作: 在执行前向或反向计算需要用到某个Layer的参数时,才通过All-Gather操作从其他GPU那里把完整的参数临时收集过来。计算完成后,非自己负责的参数就可以丢弃,释放显存。
效果: 最大程度地消除了内存冗余,使得N张卡理论上可以训练N倍大的模型(相比单卡)。
代价: 通信量大大增加,因为每次需要用到参数时都要进行All-Gather。对网络带宽要求很高。
补充:通信量增加的一些限制
一个 1B 参数的模型,每次 All-Reduce 需传输约 4GB 数据(FP32)。若网络带宽仅 10Gbps(≈1.25GB/s),仅通信就需 3.2 秒——远超前向/反向传播时间。
最直接的问题:训练速度下降;每步迭代时间(step time)增加,整体吞吐量(samples/sec)降低,这是因为梯度同步、参数广播等通信操作阻塞计算,GPU 等待网络传输完成("GPU 空转")。
同时还有扩展效率恶化;增加 GPU 数量后,加速比远低于线性(如 8 卡仅达到 4~5 倍加速),因为通信开销随节点数增长呈线性甚至超线性增长,成为瓶颈(Amdahl 定律)。
GPU 利用率波动大;通信阶段 GPU 处于等待状态,nvidia-smi 显示利用率周期性下降(如 100% → 30% → 100%),计算与通信无法充分重叠(overlap)时,资源浪费严重
CPU 与 PCIe 带宽压力;数据预处理、梯度聚合等任务占用 CPU,若 CPU 算力不足会成为新瓶颈,多卡间通过 PCIe 交换数据时,PCIe 带宽(如 x16 Gen4 ≈ 32GB/s)也可能成为限制
网络拥塞与竞争;多任务共享同一网络时(如同时跑训练+数据加载+日志上报),带宽竞争导致延迟抖动,TCP 重传、拥塞控制进一步降低有效吞吐量

ZeRO-Offload: ZeRO还可以结合CPU和NVMe(高速SSD)内存。可以将分片的优化器状态、梯度甚至参数临时“卸载(Offload)”到CPU内存或NVMe硬盘上,进一步释放宝贵的GPU显存,使得在有限的GPU资源下也能训练超大模型。当然,这会带来额外的拷贝开销,牺牲训练速度。
优点:
极大降低了数据并行下的内存冗余,能训练更大的模型或使用更大的Batch Size。
与数据并行的编程模型类似,对用户相对友好(尤其是DeepSpeed等框架封装后)。
ZeRO-3可以看作是一种动态的模型参数重组,非常灵活。
缺点:
增加了通信开销,尤其是ZeRO-3。
需要高效的通信库支持。
实现: Nccl + 高速网络硬件(IB/RoCE) 和 正确的环境变量配置(开启 RDMA)+ DeepSpeed ZeRO (Stages 1, 2, 3, Offload), PyTorch FSDP (Fully Sharded Data Parallel, 类似ZeRO-3)。
ZeRO-Offload其实提到了另一种并行方法:异构系统并行。
GPU显存不够用?把暂时不用的东西(参数、优化器状态、激活值)先挪到内存大得多的CPU或速度还行的NVMe硬盘上。
这种方式主要是为了突破硬件显存的物理限制,让你有可能在有限的GPU资源(比如单机几张卡)上,也能把一个巨大无比的模型给跑起来(哪怕慢一点)。
优点:
极大扩展了可训练模型的规模上限。
缺点:
速度慢: GPU与CPU之间(PCIe总线)、CPU与NVMe之间的数据传输速度远低于GPU之间(NVLink)或GPU内部的内存访问速度。频繁的Offload/Reload会严重拖慢训练速度。
实现复杂,需要精细管理数据的换入换出。
并行组合策略:
DP + PP:
将模型按层切分到不同的流水线阶段(PP)。
每个流水线阶段内部,再使用数据并行(DP)来处理更多数据,加速训练。
例如,你有16张卡,可以做4个流水线阶段,每个阶段内用4张卡做数据并行(4x4=16)。
DP + TP:
对于模型中特别大的层(比如Attention或MLP),使用张量并行(TP)将其切分到几张卡上。
在TP组之外,再使用数据并行(DP)来扩大训练规模。
例如,16张卡,可以做4路TP(处理大层),同时做4路DP(4x4=16)。
PP + TP:
先将模型按层切分成流水线阶段(PP)。
对于每个阶段内部比较大的层,再使用张量并行(TP)进行切分。
例如,16张卡,可以做4个流水线阶段,每个阶段内部用4张卡做TP(4x4=16)。这种组合可以处理又深又宽的模型。
DP + PP + TP (3D并行):
终极组合!同时使用数据并行、流水线并行和张量并行。
典型部署:
节点内(Intra-node): 使用TP,因为节点内GPU通过NVLink高速连接,能承受TP的大通信量。
节点间(Inter-node): 使用PP和DP,因为它们对网络带宽的要求相对较低。
例如,你有64张卡,分布在8台机器上(每台8卡)。可以:
设置TP size = 8 (节点内做张量并行)
设置PP size = 4 (跨4台机器做流水线并行)
设置DP size = 2 (剩下2组机器做数据并行)
总卡数 = 8 * 4 * 2 = 64

优点:
集各家之长,可以最大化利用集群资源,训练更大规模的模型,并获得更好的加速效果。
缺点:
极度复杂: 配置和调试非常困难,需要对各种并行策略和硬件拓扑有深入理解。
实现: DeepSpeed (支持灵活的3D并行配置), Megatron-LM (也有自己的实现), Colossal-AI, FairScale等。
MoE并行 / 专家并行:为“稀疏”而生
最后,我们简单提一下与MoE(Mixture of Experts)架构相关的并行方式。
MoE回顾: MoE模型包含多个“专家”(通常是FFN/MLP网络),和一个“门控网络(Gate)”。对于每个输入Token,门控网络会计算一个权重分布,决定将这个Token主要发送给哪些专家进行处理,最后将选定专家的输出加权组合起来。这样,模型参数量可以很大(有很多专家),但实际计算量(每个Token只激活少数几个专家)可以保持相对较低。
MoE并行(专家并行):
核心思想: 将不同的专家分布到不同的GPU上。
通信模式: 这涉及到一种All-to-All的通信。因为一个Batch中的不同Token可能会选择不同的专家,而这些专家可能位于不同的GPU上。需要将Token路由(Route)到持有对应专家的GPU上,计算完成后,再将结果收集回来。
本质: 也是一种模型并行,但其并行模式和通信模式与前面讨论的TP、PP有所不同,是专门针对MoE架构的。它通常会和数据并行结合使用(即,将整个MoE层/模型在多个GPU上做数据并行复制,同时每个副本内部的专家分布在不同的GPU上)。

优点:
允许训练参数量巨大但计算相对稀疏的MoE模型。
缺点:
All-to-All通信开销: 是主要的瓶颈,对网络带宽要求很高。
负载均衡: 如果门控网络分配不均,可能导致某些GPU上的专家很忙,而另一些很闲。
参考文章:> https://zhuanlan.zhihu.com/p/598714869
5. 大模型量化
模型量化一直以来都是将原本模型的性能在轻量化部署到弱设备端所需要进行的必要步骤,在有限的硬件资源下高效运行模型,就是本章所要解决的问题。
量化的目标,主要对象包括:
权重 (Weights): 模型参数的主体,占据了绝大部分存储空间。这是量化的首要目标。
激活 (Activations): 模型在推理过程中,层与层之间传递的中间数值。量化激活可以减少计算过程中的显存占用和访存开销。
KV Cache: 在Transformer模型的自注意力机制中,为了加速生成过程,我们会缓存过去的键(Key)和值(Value)。对于长序列,KV Cache的显存占用非常可观,对其量化也很有必要。
梯度 (Gradients) & 优化器状态 (Optimizer States): 这两者主要在训练过程中出现。量化它们可以减少训练时的显存占用和节点间的通信带宽。不过,我们今天的重点是推理优化,所以主要关注前三者
给予量化对象的不同,常见的量化方案有:
仅权重量化 (Weight-only Quantization): 比如 W4A16(权重4bit,激活16bit)、W8A16(权重8bit,激活16bit)。GPTQ 主要属于此类。
权重和激活同时量化 (Weight and Activation Quantization): 比如 W8A8(权重8bit,激活8bit)。SmoothQuant 是一个代表。LLM.int8() 在某种意义上也属于此范畴,但它更特殊一些。
KV Cache 量化: 通常采用 INT8 量化。
在介绍量化之前,我们先介绍一下为什么需要对INT8的量化做方法优化:
0.什么是INT8量化,与FP16相比有什么区别
Int8量化使用8位整数来表示模型参数和激活值,相较于传统的32位浮点数(float32)表示,其精度显著降低。这意味着在量化过程中,模型参数和激活值的取值范围被限制在一个较小的区间内,可能会导致部分模型细节的丢失,进而影响模型的最终精度。然而,这种精度损失在许多应用场景中是可以接受的,特别是在对速度要求较高、对精度要求相对较低的边缘设备或移动设备上。
INT8(8位整数) 就像一把精密游标卡尺:
• 比特分配:8位二进制数,其中1位表示符号(正负),7位表示数值
• 取值范围:-128 到 127,共256个离散值
• 精度特点:
• 相当于将权重范围分成256个"刻度"
• 每刻度间的跳跃相对平缓,误差易于控制
• 能够较精确地表示大多数权重值
其核心原理主要包括数据缩放与取整两个步骤。数据缩放需要确定缩放因子和零点偏移,以将原始浮点数据的范围映射到int8的表示范围内,int8的表示范围通常为-128到127。常见的映射方法包括对称量化和非对称量化。随后,将缩放后的浮点数据四舍五入到最接近的整数值,完成量化。
per_channel_quant_int8
按通道 INT8 量化是一种将浮点数张量(通常是 FP32 或 FP16)压缩为 INT8 的量化方法,其中每个通道使用独立的量化参数(scale 和 zero-point)。 在深度学习模型中,例如线性层或卷积层的权重矩阵,不同通道的数据分布可能差异很大。
某些通道数值范围可能在 [-0.2, 0.2]。某些通道数值范围可能在 [-10, 10]。如果使用统一的量化参数(Per-Tensor Quantization):小范围通道会严重损失精度,大范围通道可能发生溢出。
因此采用 Per-Channel Quantization,即每个通道独立计算:scale(缩放因子),zero-point(零点)。
例如对于一个形状为x.shape = [M, K] 如果按 dim=-1 做 per-channel 量化:每一行共享一个 scale, scale 的形状为[M,1]。这样每一行可以使用不同的量化比例,从而显著提高量化精度。
点击查看代码
def per_channel_quant_int8_torch(x, symmetric):
if symmetric:
x = x.float()
absmax = x.abs().max(dim=-1).values
absmax = absmax.clamp_min(1e-10).unsqueeze(-1)
scale_x = absmax / 127
x_q = x.mul(127 / absmax)
x_q = torch.round(x_q).to(torch.int8)
return x_q, scale_x, None
else:
w = x.float()
w_min = w.min(dim=-1, keepdim=True)[0]
w_max = w.max(dim=-1, keepdim=True)[0]
w_scale = (w_max - w_min) / 255.0
w_scale = torch.clamp(w_scale, min=1e-8)
w_zero = -w_min / w_scale - 128.0
w_q = torch.round(w / w_scale + w_zero)
w_q = torch.clamp(w_q, -128, 127)
w_packed = w_q.to(torch.int8)
return w_packed, w_scale, w_zero
对称量化的特点:
zero-point 为 0;数值范围为 [-127,127];计算更简单;更适合硬件加速
非对称量化的特点:
zero-point 不为 0
数值范围为 [-128,127]
可以更好适配非对称分布数据
计算稍微复杂 因此在很多推理框架中,例如:TensorRT、CUTLASS、AWQ、GPTQ、Marlin 通常更倾向于使用对称量化,因为在矩阵乘法(INT8 GEMM)中,zero-point 为 0 可以减少额外的计算。
per_tensor_quant_int8
逐张量 INT8 量化是一种常见的量化方式,其核心思想是:整个张量共享同一个量化参数(scale),通过一个统一的缩放比例将浮点数据映射到 INT8 表示范围内。
点击查看代码
def per_tensor_quant_int8_torch(x, symmetric):
if symmetric == False:
return
else:
x = x.float()
absmax = x.flatten().abs().max()
if absmax == 0:
scale = torch.tensor(1.0, device=x.device, dtype=torch.float32)
q = torch.zeros_like(x, dtype=torch.int8)
return q, scale, None
scale_x = absmax / 127
x_q = x.mul(127 / absmax)
x_q = torch.round(x_q).to(torch.int8)
return x_q, scale_x, None
在 Per-Tensor 量化中,首先需要计算整个张量中元素的最大绝对值: absmax = max(|x|)
然后根据该最大值计算量化比例(scale): scale = absmax / 127
这个 scale 的含义是:将浮点数范围 [-absmax, absmax] 线性映射到 INT8 范围 [-127,127]。
随后可以通过如下公式完成量化: q = round(x / scale) 或者等价写法: q = round(x * 127 / absmax) 量化后的 q 即为 INT8 表示的数据。
若需要恢复近似的浮点值,则可以通过反量化公式: x ≈ q * scale 来获得。
在给出的代码实现中,函数首先将输入张量转换为 float 类型,以确保后续计算精度。然后通过: absmax = x.flatten().abs().max() 计算整个张量的最大绝对值。这里通过 flatten() 将张量展平,是为了在所有维度上统一计算最大值,从而得到全局量化范围。 如果张量所有元素都为 0,则 absmax 会等于 0,此时 scale 会变为 0 并导致除零错误。
因此代码中特别处理了这种情况:当 absmax == 0 时,直接返回全 0 的 INT8 张量,并将 scale 设置为 1.0。 否则按照正常流程计算: scale = absmax / 127 并执行量化: x_q = round(x * 127 / absmax) 最终得到 INT8 类型的量化结果 x_q 以及对应的 scale。
由于这里采用的是对称量化(symmetric quantization),zero-point 始终为 0,因此函数返回的 zero-point 为 None。 Per-Tensor Quantization 与 Per-Channel Quantization 的主要区别在于量化参数的共享方式。
在 Per-Tensor Quantization 中:
整个张量共享一个 scale;量化参数数量少;实现简单;计算开销低。
例如对于一个矩阵:x.shape = [M, K]
Per-Tensor量化只会计算:scale : 标量。整个矩阵所有元素都会使用同一个 scale 进行量化。
而在 Per-Channel Quantization 中:每个channel使用独立的scale,不同通道可以适应不同的数据分布,量化精度更高。
例如对于同样的矩阵: x.shape = [M, K]。
如果按dim=-1做 per-channel 量化,则每一行会拥有独立的scale:scale.shape = [M, 1],量化计算时会利用 broadcast 机制,使每一行的数据按照不同的 scale 进行量化。

在深度学习推理中,通常会采用:
Per-Tensor Quantization 用于激活(activation)
Per-Channel Quantization 用于权重(weight) 原因是权重在不同通道上的数值分布差异往往较大,如果使用统一 scale,可能会造成明显的精度损失;而 Per-Channel 量化可以显著缓解这一问题,因此在现代推理框架(如 TensorRT、CUTLASS、AWQ、GPTQ 等)中被广泛使用。
1. 零样本量化(不需要额外的校准数据集):LLM.int8()
在模型的激活值中,存在一些数值特别大(绝对值远超其他值)的“离群值”(Outliers)。更关键的是,这些离群值往往集中在少数几个特征维度(Feature Dimensions)上。
以一个典型的矩阵乘法 Y = XW 为例,X 是输入激活(形状 [序列长度 T, 隐藏层维度 h]),W 是权重(形状 [隐藏层维度 h, 输出维度 h0])。这里的“特征维度”主要指 h 这个维度。
传统的量化方法,无论是 per-token(对 X 的每一行使用一个量化缩放因子)还是 per-channel(对 W 的每一列使用一个量化缩放因子),都很容易被这些极端离群值“带偏”。为了容纳这些离群值,量化范围必须拉得很大,导致大部分正常数值的精度损失严重,就像用大盒子装小苹果一样。
那LLM.int8()的想法就是,将离群值的特征维度单独拿出来处理,其余部分进行常规的INT8量化处理,即混合精度分解。
对于矩阵乘法 Y = XW,LLM.int8() 的计算过程大致分为三步:
1. 识别与分离 (Identify & Separate):
首先,确定一个阈值(比如 6.0)。
遍历输入激活 X 的每一列(即每个特征维度)。如果某一列包含绝对值大于阈值的离群值,则该列对应的整个特征维度被标记为“离群特征维度”。
将输入 X 分解为两部分:X_outlier(只包含离群特征维度的列,其余列为0)和 X_regular(只包含非离群特征维度的列,其余列为0)。
相应地,权重 W 也根据“离群特征维度”的行(因为 W 的行对应 X 的列)进行分解,得到 W_outlier 和 W_regular。
2.分别计算 (Separate Computation):
离群部分: 对离群特征维度,执行高精度(FP16)的矩阵乘法:Y_outlier = X_outlier @ W (或者等价地 X @ W_outlier)。注意,这里虽然分离了 X_outlier,但乘法通常还是和完整的 W (或者X与W_outlier) 进行,只是因为 X_outlier 中很多列是0,实际计算量集中在离群维度上。 (注:原论文图示更侧重分解 W,但实际操作中分解 X 更常见,核心思想一致)
常规部分: 对非离群特征维度,执行 INT8 量化计算:
将 X_regular 进行 vector-wise (per-token) INT8 量化,得到 X_regular_int8 和量化尺度 S_X。
将 W_regular 进行 row-wise/column-wise (per-channel) INT8 量化,得到 W_regular_int8 和量化尺度 S_W。
执行 INT8 矩阵乘法:Y_regular_int8 = X_regular_int8 @ W_regular_int8。
3.合并结果 (Combine Results):
将 INT8 乘法结果反量化回 FP16:Y_regular_fp16 = Dequantize(Y_regular_int8, S_X, S_W)。
将离群部分和常规部分的结果相加,得到最终的 FP16 输出:Y = Y_outlier + Y_regular_fp16。
这个方法为什么能work,基于一个关键观察:
离群特征的出现并非模型规模独有,而是与模型的困惑度(perplexity)更相关。随着模型训练得更好(困惑度降低),离群特征会更显著。
大约在 6B 参数规模附近,离群特征出现的普遍性(影响的层数和特征维度比例)和强度(离群值的数值大小)会有一个“相变”式的急剧增长。这恰好解释了为什么普通 INT8 量化在这个规模开始失效。

尽管该量化解决了常规int8给模型带来的性能损失,但是,这种混合精度计算和条件判断(分离离群值)引入了额外的开销,导致推理速度相较于纯 FP16 反而变慢了。对于 BLOOM-176B,速度大约降低了 15%-23%;对于小模型,这个差距可能更大。所以,LLM.int8() 主要的优势在于显著降低显存占用,使得大模型能在消费级硬件上运行起来,而不是追求极致的推理速度。
点击查看代码
import torch
from transformers import AutoModelForCausalLM
# 假设你有足够的总显存,但单卡可能不够
# device_map='auto' 会自动将模型分片到多张GPU上
# load_in_8bit=True 启用 LLM.int8() 量化
model_8bit = AutoModelForCausalLM.from_pretrained(
'bigscience/bloom-1b7', # 换成你想用的模型
device_map='auto',
load_in_8bit=True,
# 可选:为每张卡设置最大显存限制
# max_memory={i: f'{int(torch.cuda.mem_get_info(i)[0]/1024**3)-2}GB' for i in range(torch.cuda.device_count())}
)
print(f"模型显存占用: {model_8bit.get_memory_footprint() / 1e9:.2f} GB")
2. 基于优化的精准权重量化:GPTQ
与 LLM.int8() 不同,GPTQ 是一种仅权重量化(Weight-only Quantization)方法,通常采用 W4A16(权重4bit,激活FP16)的方案。它不处理激活值的离群问题,而是专注于如何更精确地量化权重,以最小化量化带来的误差。
GPTQ 的思想并非凭空产生,它借鉴了经典的模型剪枝方法。
OBD (Optimal Brain Damage): 早期剪枝方法,假设移除一个权重对模型误差的影响可以通过该权重对应的 Hessian 矩阵(二阶导数)的对角线元素来近似。它假设权重间影响独立。
OBS (Optimal Brain Surgeon): 改进了 OBD,认为权重间并非独立。它不仅计算移除哪个权重影响最小,还计算如何调整剩余权重以补偿移除权重带来的误差。这需要计算 Hessian 矩阵的逆。
OBQ (Optimal Brain Quantization): 将 OBS 的思想从“剪枝”(将权重置为0)推广到“量化”(将权重近似到某个离散值)。量化可以看作是找到一个量化后的权重\(\hat{W}\),使得\(||WX - \hat{W} X||^2\)尽可能小,同时考虑了更新其他权重来补偿误差。
OBQ 的效果不错,但计算 Hessian 矩阵的逆非常耗时,对于动辄百亿、千亿参数的大模型来说,计算成本高到无法接受。
GPTQ (Generative Pre-trained Transformer Quantization) 的目标就是:在继承 OBQ 精确量化思想的同时,大幅提高量化速度,使其能够应用于超大规模模型。
技术原理:逐层优化与权重更新
GPTQ 采用逐层量化 (Layer-wise Quantization) 的策略,即一次只量化模型的一层。对于每一层,它的目标是找到量化后的权重\(\hat{W}\),使得该层的输出与原始输出之间的均方误差最小:\(argmin || WX - round(W)X ||^2\)
这里 X XX 是该层的输入激活(通常来自一个小的校准数据集),\(\text{round}(\cdot)round(⋅)\) 代表量化操作。
GPTQ 的核心是逐列(或逐块)量化权重矩阵W,并实时更新还未量化的权重,以补偿已量化权重引入的误差。这个过程有点像 OBS,但 GPTQ 做了几项关键优化:
按顺序量化,而非贪心选择: OBS/OBQ 需要计算 Hessian 逆来决定量化哪个权重影响最小(贪心策略)。GPTQ 发现,直接按顺序(例如,从第一列到最后一列)量化权重,性能损失很小。这大大简化了计算,避免了复杂的排序和索引。
分组量化 (Group-wise Quantization): 为了进一步平衡精度和效率,GPTQ 不是逐个元素量化,而是将权重矩阵的列(或行,取决于实现)分成小组(例如,每组 128 列,称为 groupsize。组内的权重共享相同的量化参数(缩放因子和零点)。这降低了存储量化参数的开销。
权重更新: 这是关键步骤。当一个权重(或一个 group 的权重)\(w_q\)被量化后,它会引入一个误差 \(error = w_q - \text{round}(w_q)\))。GPTQ 利用预计算的 Hessian 逆的块信息 (\(H^{-1}\)),将这个误差分配到还未被量化的其余权重 \(W_{\text{remaining}}\)上:

高效实现技巧:
Lazy Batch-Updates: OBQ 每次量化一个权重就更新所有剩余权重,这在 GPU 上访存效率低。GPTQ 采用延迟批量更新,一次量化一个小批量(如 128 列)的权重,然后一次性更新所有剩余权重。这能更好地利用 GPU 的并行计算能力和内存带宽。
Cholesky 分解: 直接计算和存储 Hessian 矩阵的逆\(H^{-1}\)可能数值不稳定且耗内存。GPTQ 使用 Cholesky 分解等数值稳定的方法来处理 Hessian 矩阵,并在需要时高效地计算更新所需的项,提高了稳定性和效率。
虽然 GPTQ 降低了时间复杂度,但这个算法的计算/通信比太低,通信带宽成为了瓶颈。
例如在量化某个参数矩阵的情况下,每次量化一个参数,其他所有未量化的参数都要按公式全都要更新一遍:
如果每行的量化并行计算,那么每次更新过程就需要 read + write 一次参数矩阵。如果参数矩阵的维度为k * k,那么量化这个参数矩阵就需要读写 k 次参数,总共的 IO 量为k^3个元素。当 k 比较大时(>= 4096),需要读写的元素就非常多了,运行时间大都被 IO 占据。
思路:由于参数量化是一列一列按次序进行的,第 i 列的参数的量化结果受到前 i-1 列量化的影响,但第 i 列的量化结果不影响前面列的量化。因此我们不需要每次量化前面的列,就更新一遍第 i 列的参数,而是可以先记录第 i 列的更新量,在量化到第 i 列时,再一次性更新参数,这样就可以减少 IO 的次数。
具体实现:将参数矩阵按每 128 列划分为一个个 group,量化某一列时,group 内的参数立即更新,而 group 后面的列只记录更新量,延迟更新。当一个 group 的参数全部量化完成,再统一对后面的所有参数做一次更新。这就是 Lazy Batch-Updates 。

Lazy Batch-Updates 不减少实际的计算量,但它能有效解决吞吐的瓶颈问题。
优点:
高精度: 由于其基于优化的量化策略和权重更新机制,GPTQ 通常能在极低比特(如 4-bit)下保持相当高的模型精度,性能损失很小。
显著的内存节省: W4A16 方案能将模型权重大小压缩近 4 倍。
潜在的推理加速: 虽然激活仍然是 FP16,但权重的位宽降低,可以减少访存时间,尤其是在内存带宽受限的场景下,可能带来推理速度的提升(但这不绝对,需要特定的 Kernel 优化支持)。
缺点:
需要校准数据: GPTQ 需要一小部分代表性的数据(校准集)来计算激活 X 和 Hessian 矩阵 H,用于指导量化过程。
量化过程耗时: 虽然比 OBQ 快得多(数小时 vs 数天),但相比 LLM.int8() 的零样本加载还是要慢不少,需要在部署前进行离线量化。
实现与应用:
目前社区有多个流行的 GPTQ 实现库:
AutoGPTQ: 一个用户友好、支持多种 Transformer 架构的库,并且已经集成到 Hugging Face 的 transformers 和 optimum 库中,使用非常方便。
GPTQ-for-LLaMa, Exllama, llama.cpp: 这些库通常针对 LLaMA 及其变种模型做了深度优化,可能在特定模型上性能更好,但通用性不如 AutoGPTQ。
3. SmoothQuant: 实现高精度且高效率的W8A8:
针对异常值的问题,一般采用更细粒度的量化方法:
逐层量化 (Per-tensor): 整个权重或激活张量共享一对量化参数 (scale s 和 zero-point z)。计算效率最高,但精度损失可能最大,尤其在有异常值时。
逐通道量化 (Per-channel / Per-token):
Per-channel (通常用于权重): 权重的每一列有自己的量化参数。
Per-token (通常用于激活): 激活的每一个 Token (行) 有自己的量化参数。
Vector-wise 是 Per-token/Per-channel 的一种统称或特定实现。
这种方式能更好地适应数据的局部变化,精度较高。
逐组量化 (Per-group / Group-wise): 将通道分成小组,组内共享量化参数。是 Per-tensor 和 Per-channel 的折中。
研究发现,对于激活量化,只有逐通道 (per-channel) 的方式(即激活的每个特征维度/列有独立的量化参数)能够较好地保持精度。但是,这种方式与现代硬件(如 GPU)上的高效INT8 GEMM (通用矩阵乘法) 内核不兼容!硬件上的 INT8 加速单元通常是为 per-tensor 或 per-token 量化设计的,无法高效执行 per-channel 的激活量化,导致推理速度上不去。
为了进行 vector-wise quantization 以有效利用 INT8 GEMM 内核,只能使用外部维度(即Token维度 T 和 通道外维度 C0)的缩放因子,不能使用内部维度(即通道内维度 Ci),因此,先前的工作对激活都采用了per-token量化,但并不能降低激活的难度。
LLM.int8使用混合精度避开了这个问题,但牺牲了速度,SmoothQuant的想法就是,将激活量化的难度转移一部分到权重上。

这个转换的结果就是,激活值X的异常值被削峰,分布更加均匀,权重虽然出现了些峰值,但相比激活值来说量化还是更容易的。那么,如何确定最佳的平滑因子?

这样来看,取0.5来说就是最均衡的。对于异常值突出的模型,用更大的参数会更好一些。
SmoothQuant工作流程分为两个阶段:离线处理和在线推理两个阶段。






4. AWQ:激活感知权重量化
量化如果进一步到INT4,INT3,一般用两种方法:
量化感知训练 (Quantization-Aware Training, QAT): 虽然效果好,能在训练中模拟量化误差并进行调整,但需要完整的训练或微调流程,计算成本高昂,对于动辄千亿参数的 LLM 来说不太现实。
训练后量化 (Post-Training Quantization, PTQ): 无需重新训练,成本低廉。但传统的 PTQ 方法(如简单的取整 RTN)在低比特场景下精度损失严重。像 GPTQ 这样的先进 PTQ 方法,虽然利用二阶信息进行误差补偿,提高了精度,但它在量化过程中可能“用力过猛”,过度拟合用于校准的小数据集,损害了 LLM 作为“通才模型”在广泛任务上的泛化能力。
我们希望找到一种PTQ方法,既能实现低比特权重量化,保持高精度,又不损害模型的泛化能力,还对硬件友好。这就是AWQ技术。
AWQ 的出发点非常直观:LLM 中的权重并非生而平等,并非所有权重都同等重要。 存在一小部分(约 0.1% - 1%)的显著权重 (salient weights),它们对模型的整体性能起着决定性的作用。
AWQ的第一个核心观察:权重的显著性是由激活感知的 (Activation-aware)。因为激活值代表了流经模型的信息。那些经常处理具有较大激活幅度的特征的权重通道,自然承载了更重要的信息流。 保留这些与大激活值对应的权重通道(保持 FP16),效果立竿见影!仅仅保留0.1% - 1%的这类权重,量化后的模型性能就得到了显著提升,甚至能媲美复杂的 GPTQ 方法。
AWQ的第二个核心参看:SmoothQuant,激活感知缩放。保护这些重要的权重,又不引入硬件不友好的混合精度。

如何找到最优的缩放因子s?


最后一个核心:权重裁剪,量化之前,将权重中超过某个阈值的数值裁剪掉,进一步降低量化权重的最大值,减小量化步长,降低整体量化误差。

AWQ 的优势:
高精度: 在 INT4/INT3 权重量化下,精度显著优于 RTN,与 GPTQ 相当或更好。
强泛化能力: 不依赖反向传播或复杂的重构,仅使用激活尺度信息,对校准集依赖小,不易过拟合,能很好地保持 LLM 在各种任务上的原始能力。
硬件友好: 最终产物是纯粹的低比特权重,没有混合精度问题,易于在现有硬件上实现高效推理(通常是 W4A16 或 W3A16 形式,即 INT4/3 权重 + FP16 激活/计算)。
速度快: 量化过程本身非常快(主要是网格搜索 α)。
AutoAWQ量化LLM示例:
点击查看代码
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
model_path = "facebook/opt-125m" # 你想量化的模型
quant_path = "opt-125m-awq" # 量化后模型的保存路径
# 定义量化配置
quant_config = {
"zero_point": True, # 是否使用零点
"q_group_size": 128, # 权重量化的组大小
"w_bit": 4, # 权重比特数 (4-bit)
"version": "GEMM" # AWQ kernel 版本 (GEMM/GEMV)
}
# 加载模型和 Tokenizer
model = AutoAWQForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True) # 使用 low_cpu_mem_usage 减少内存占用
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
# 执行量化 (需要少量校准数据,这里省略了数据加载步骤,AutoAWQ 默认会使用 wiktext 或 c4)
print("Starting quantization...")
model.quantize(tokenizer, quant_config=quant_config)
print("Quantization complete.")
# (注意:实际量化需要一些时间,具体取决于模型大小和硬件)
<details>
<summary>点击查看代码</summary>
from transformers import AwqConfig
创建符合 Transformers 格式的 AwqConfig
quantization_config = AwqConfig(
bits=quant_config["w_bit"],
group_size=quant_config["q_group_size"],
zero_point=quant_config["zero_point"],
version=quant_config["version"].lower(),
)
将量化配置附加到模型配置中
model.model.config.quantization_config = quantization_config
保存量化后的模型权重和 Tokenizer
print(f"Saving quantized model to {quant_path}...")
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
print("Model saved.")
</details>

目前来看,AWQ(及其代表的 W4A16/W3A16 这类仅权重量化方案)在社区和实际应用中似乎更为主流和受关注,尤其是在大型语言模型的部署场景下。
主要原因如下:
内存是首要瓶颈: 对于非常大的 LLM(几十B到几百B参数),显存容量 (VRAM) 往往是部署的最大障碍。将权重从 FP16 量化到 INT4/INT3 可以直接将模型大小减少 75%-80%!这种数量级的减少使得原本无法在单张甚至多张消费级/专业级 GPU 上运行的模型成为了可能。相比之下,SmoothQuant (W8A8) 只能减少 50% 的内存占用,吸引力相对较小。
W4A16 精度足够好: 令人惊讶的是,像 AWQ 和 GPTQ 这样的方法证明了,即使只用 4 比特量化权重,只要方法得当(比如保护关键权重),LLM 仍然可以在很多任务上保持相当高的精度,损失在可接受范围内。这使得 W4A16 成为了一个性价比极高的选择。
推理速度瓶颈的变化: 在很多实际的 LLM 推理场景(尤其是小批量或单用户交互),内存带宽 (Memory Bandwidth) 往往是比计算本身更大的瓶颈(Memory-bound)。读取更小的 INT4/INT3 权重可以显著降低对内存带宽的需求,从而带来实际的推理加速,即使计算本身仍然是 FP16。SmoothQuant 虽然利用了 INT8 计算,但在内存带宽受限场景下,其加速效果可能不如 W4A16 带来的内存读取优势明显。
浙公网安备 33010602011771号