目录
- Abstract
- 1 Introduction
- 2 BitNet
- 2.1
- 2.2 模型训练
- 2.3 计算效率
- 3 与 FP16 Transformers 的比较
- 3.1 设置
- 3.2 推理优化的缩放法则
- 3.3 下游任务的结果
- 3.4 稳定性测试
- 4 与训练后量化的比较
- 4.1 设置
- 4.2 结果
- 5 消融研究
- 6 结论与未来工作
Abstract
大型语言模型的规模不断增加,给部署带来了挑战,并因高能耗而引发了对环境影响的担忧。在这项工作中,我们介绍了BitNet,一种为大型语言模型设计的可扩展且稳定的1位Transformer架构。具体来说,我们引入了BitLinear,作为nn.Linear层的替代品,以从头开始训练1位权重。语言建模的实验结果表明,BitNet在显著减少内存占用和能耗的同时,实现了与最先进的8位量化方法和FP16 Transformer基线相竞争的性能。此外,BitNet表现出类似于全精度Transformers的缩放规律,这表明它在保持效率和性能优势的同时,有潜力有效扩展到更大的语言模型。
1 Introduction
大型语言模型的快速增长 [BMR+20, Ope23, CND+22, ADF+23, TLI+23, TMS+23] 在各种任务中取得了显著的改进。然而,由于高昂的推理成本和能耗,托管大型语言模型的费用非常高。随着这些模型规模的扩大,访问和处理模型参数所需的内存带宽成为主要瓶颈,限制了整体推理性能。此外,当在分布式系统或多设备平台上部署这些模型时,设备间通信的开销会显著影响推理延迟和能量消耗。模型量化 [FAHA23, CCKS23, XLS+23] 已经成为一种有前景的解决方案,因为它可以显著减少大规模模型的内存占用和计算成本,同时保持竞争力的性能。
大多数现有的大型语言模型的量化方法都是训练后量化。它们简单易用,因为不需要对训练流程进行任何修改或重新训练模型。然而,当精度降低时,该方法会导致更显著的准确性损失,因为模型在训练期间未针对量化表示进行优化。
另一种量化深度神经网络的方法是量化感知训练。与训练后量化相比,它通常会带来更好的准确性,因为模型从一开始就考虑了精度降低。此外,它允许模型继续训练或进行微调,这对大型语言模型尤为重要。量化感知训练的挑战主要在于优化,即随着精度的降低,模型变得更难收敛。此外,尚不清楚量化感知训练是否遵循神经语言模型的缩放规律。
在这项工作中,我们专注于二值化(即1位),这是量化的极端情况,应用于大型语言模型。以前关于二值化神经网络的研究 [RORF16, BT19] 主要集中在卷积神经网络上。最近,有一些关于二值化Transformers的研究。然而,这些研究主要集中在机器翻译或BERT预训练上,与大型语言模型有很大不同。例如,机器翻译采用编码器-解码器架构,BERT预训练使用双向编码器,而大型语言模型使用单向解码器。此外,大型语言模型通常扩展到更大的模型规模,而BERT和机器翻译模型并没有经历如此广泛的扩展。
据我们所知,这项工作是首次研究1位大型语言模型的量化感知训练。我们提出了BitNet,一种用于大型语言模型的1位Transformer架构,旨在在内存和计算方面高效扩展。BitNet在训练期间使用低精度的二进制权重和量化激活,同时保持优化器状态和梯度的高精度。我们的方法设计为可扩展且稳定,能够高效处理大型语言模型。BitNet架构的实现非常简单,只需替换Transformer中的线性投影(即PyTorch中的nn.Linear)。此外,它还可以与其他大型语言模型的加速方法互补,如PagedAttention [KLZ+23]、FlashAttention [DFE+22, Dao23] 和推测解码 [LKM23]。
我们在一系列语言建模基准上评估了BitNet,与最先进的量化方法和FP16 Transformers进行比较。实验结果表明,BitNet在困惑度和下游任务准确性方面表现出竞争力。更重要的是,BitNet显著减少了内存占用和能量消耗。此外,我们表明BitNet遵循类似于全精度Transformers的缩放规律,表明它可以有效扩展到更大的语言模型,并在性能和效率方面具有潜在优势。
2 BitNet
如图2所示,BitNet使用与Transformers相同的布局,堆叠自注意力机制和前馈网络的块。与原始Transformer相比,BitNet使用BitLinear(公式11)代替传统的矩阵乘法,后者采用二值化(即1-bit)的模型权重。我们将其他组件保持高精度,例如在实验中使用8-bit。我们总结了以下原因。首先,残差连接和层归一化对大型语言模型的计算成本贡献可以忽略不计。其次,随着模型的增大,QKV转换的计算成本远小于参数投影的计算成本。第三,我们保留输入/输出嵌入的精度,因为语言模型必须使用高精度的概率来进行采样。
2.1
2.1 BitLinear
我们首先使用符号函数将权重二值化为+1或-1。根据[LOP+22]的方法,在二值化之前,我们将权重中心化为零均值,以在有限的数值范围内增加容量。在二值化之后使用一个缩放因子β,以减少实值权重和二值化权重之间的l2误差。一个权重矩阵W ∈ Rn×m的二值化可以表示为:
[ W = \text{Sign}(W - α), ]
[ \text{Sign}(W_{ij}) = \begin{cases}
+1, & \text{if } W_{ij} > 0 \
-1, & \text{if } W_{ij} \leq 0
\end{cases} ]
[ α = \frac{1}{nm} \sum_{ij} W_{ij} ]
我们进一步将激活值量化为b-bit精度。根据[DLBZ22]的方法,我们使用absmax量化,它通过将激活值乘以Qb并除以输入矩阵的绝对最大值,将激活值缩放到范围[-Qb, Qb] (Qb = 2^b−1):
[ x = \text{Quant}(x) = \text{Clip}\left(\frac{x \times Qb}{γ}, -Qb + ϵ, Qb - ϵ \right), ]
[ \text{Clip}(x, a, b) = \max(a, \min(b, x)), ]
[ γ = ||x||_∞, ]
其中ϵ是防止溢出的小浮点数。
对于非线性函数(例如ReLU)之前的激活值,我们通过减去输入值的最小值将它们缩放到范围[0, Qb],以确保所有值为非负:
[ x = \text{Quant}(x) = \text{Clip}\left(\frac{(x - η) \times Qb}{γ}, ϵ, Qb - ϵ \right), ]
[ η = \min_{ij} x_{ij}. ]
在本工作中,我们将激活值量化为8-bit,并在未来工作中探索更低的精度。此外,训练期间按张量进行量化,而推理期间按token进行量化,以确保稳定性和效率。
根据上述量化方程,矩阵乘法可以表示为:
[ y = Wx ]
我们假设W和x中的元素相互独立并共享相同的分布,并且W和x彼此独立。那么输出y的方差估计为:
[ \text{Var}(y) = n \text{Var}(wx) ]
[ = n E[w^2] E[x^2] ]
[ = n β^2 E[x^2] ≈ E[x^2] ]
对于全精度计算,输出的方差Var(y)在使用标准初始化方法(例如Kaiming初始化或Xavier初始化)时通常为1,这对训练稳定性有很大好处。为了在量化后保持方差,我们在激活量化之前引入了LayerNorm [BKH16]函数。通过这种方式,输出y的方差估计为Var(y) ≈ E[LN(x)^2] = 1,这与全精度的Var(y)具有相同的量级。在Transformers的上下文中,它与SubLN [WMH+22]具有相同的实现。使用SubLN和上述量化方法,我们得到BitLinear,其公式为:
[ y = Wx = W \text{Quant}(\text{LN}(x)) \times \frac{βγ}{Qb} ]
[ \text{LN}(x) = \frac{x - E(x)}{\sqrt{\text{Var}(x)} + ϵ}, β = \frac{1}{nm ||W||_1} ]
图2展示了BitLinear的计算流程。在SubLN操作之后,激活值使用absmax函数进行量化。然后在1-bit权重和量化激活值之间进行矩阵乘法。输出激活值通过{β,γ}重新缩放,以将它们反量化到原始精度。
模型并行化与组量化和归一化
扩展大型语言模型的一个基本技术是模型并行化[SPP+19],它将矩阵乘法分割到多个设备上。现有模型并行化方法的前提是张量在分区维度上是独立的。然而,所有参数α, β, γ和η都是从整个张量计算出来的,破坏了独立性前提。一个解决方案是为每个参数引入一个all-reduce操作。然而,即使每个参数的通信量很小,随着模型变得更深,同步的数量也在增加,这显著减慢了前向传递。这个问题也存在于SubLN中,其中均值和方差需要在分区维度上估计。
为此,我们提出了一种使模型并行化更高效的简单方法。我们将权重和激活值划分为多个组,然后独立估计每个组的参数。这样,参数可以在本地计算,而无需额外的通信。这种方法称为组量化,其公式如下:
对于一个权重矩阵W ∈ Rn×m,我们沿着分区维度将其分为G组,每组大小为(\frac{n}{G} × m)。然后我们独立估计每组的参数:
[ α_g = \frac{G}{nm} \sum_{ij} W(g)_{ij}, β_g = \frac{G}{nm ||W(g)||_1} ]
其中W(g)表示权重矩阵的第g组。同样,对于激活值,我们可以将输入矩阵x∈Rn×m分为G组,并计算每组的参数:
[ γ_g = ||x(g)||∞, η_g = \min{ij} x(g)_{ij} ]
对于LN,我们可以应用组归一化技术[WH20],独立计算每组的均值和方差:
[ \text{LN}(x(g)) = \frac{x(g) - E(x(g))}{\sqrt{\text{Var}(x(g))} + ϵ} ]
通过这种方式,我们可以高效地实现具有组量化和归一化的模型并行化,无需额外的通信,并且可以扩展到大型语言模型。
2.2 模型训练
直接估计器
为了训练我们的1-bit模型,我们采用了直接估计器(STE)[BLC13]在反向传播过程中近似梯度。这种方法在反向传递过程中绕过了不可微函数,如符号函数(公式2)和剪辑函数(公式5)。STE允许梯度在网络中流动而不受这些不可微函数的影响,从而使得训练我们的量化模型成为可能。
混合精度训练
虽然权重和激活值被量化到低精度,但梯度和优化器状态被存储为高精度,以确保训练的稳定性和准确性。根据之前的工作[LSL+21],我们维护了一个高精度格式的潜在权重,以便累积可学习参数的更新。潜在权重会在前向传递过程中动态二值化,并且永远不会用于推理过程。
大学习率
优化的一大挑战在于,潜在权重上的小更新常常对1-bit权重没有影响。这会导致基于1-bit权重估计的有偏梯度和更新。在训练开始阶段,这个问题尤为严重,因为模型需要尽快收敛。为了解决这个问题,我们探索了各种方法,最后得出结论增加学习率是加速优化的最简单和最有效的方法。我们的实验表明,BitNet在收敛速度方面受益于大学习率,而FP16 Transformer在训练开始时使用相同的学习率会发散。更多细节可以在第3节中找到。
2.3 计算效率
我们从算术操作能量和内存占用的角度估计了BitNet的计算效率。我们主要关注矩阵乘法的计算,因为它对大型语言模型的成本贡献最大。
算术操作能量
根据[Hor14, ZZL22]中的能量模型,不同算术操作的能量消耗可以估算如下:
位数 | 加法能量估算 ( \hat{E}_{add} ) (pJ) | 乘法能量估算 ( \hat{E}_{mul} ) (pJ) |
---|---|---|
45nm | 7nm | |
FP32 | 0.9 | 0.4 |
FP16 | 0.03 | 0.16 |
INT8 | 0.007 | 0.38 |
表2:不同位数表示在45nm和7nm工艺节点上的加法和乘法能量消耗[Hor14, ZZL22]。
在普通Transformers中,对于维度为 (m \times n) 和 (n \times p) 的矩阵乘法,能量消耗可以计算如下:
[ E_{add} = m \times (n - 1) \times p \times \hat{E}_{add} ]
[ E_{mul} = m \times n \times p \times \hat{E}_{mul} ]
对于BitNet,由于权重是1-bit,矩阵乘法的能量消耗主要由加法操作主导。乘法操作只用于用标量 (\beta) 和 (\frac{\gamma}{Q_b}) 缩放输出,所以乘法的能量消耗可以计算如下:
[ E_{mul} = (m \times p + m \times n) \times \hat{E}_{mul} ]
这比Transformers中的能量消耗要小得多。表1展示了W1A8 BitNet与全精度(32-32)和半精度(16-16)Transformer相比的能量节省情况。可以看到,BitNet提供了显著的能量节省,特别是在乘法操作方面,这是矩阵乘法能量消耗的主要组成部分。
3 与 FP16 Transformers 的比较
3.1 设置
我们训练了一系列不同规模的自回归语言模型,BitNet的参数量从125M到30B不等。模型在一个英语语料库上进行训练,该语料库包括Pile数据集、CommonCrawl快照、RealNews和CC-Stories数据集。我们使用SentencePiece分词器对数据进行预处理,词汇表大小为16K。除了BitNet,我们还使用相同的数据集和设置训练了Transformer基线模型,以进行公平的比较。更多细节可以在附录中找到。
3.2 推理优化的缩放法则
神经语言模型已经证明可以按照可预测的方式扩展[KMH+20],特别是在普通Transformer架构下。损失与用于训练的计算量成幂律关系。这使我们能够确定计算预算的最佳分配,并从较小的模型中预测大型语言模型的性能。
为了研究二值化Transformer的缩放法则,我们首先绘制了BitNet和FP16 Transformer基线模型的缩放曲线,比较它们的参数量。我们固定训练token的数量并改变模型的大小。图3显示,BitNet的损失缩放与FP16 Transformer相似,都遵循幂律。我们随后使用一个不可约损失项来拟合缩放法则:
[ L(N) = aN^b + c ]
为了评估缩放法则是否能够准确预测损失,我们选择了从125M到6.7B的模型来拟合幂律中的参数,并使用该法则来预测13B和30B模型的损失。结果表明,拟合的缩放法则能够高度准确地预测BitNet的损失。此外,随着模型规模的增长,BitNet和FP16 Transformer之间的差距变得更小。
虽然上述幂律度量了BitNet的缩放趋势,但它并不能准确地建模损失和实际计算量之间的关系。之前的工作 [KMH+20, HKK+20, HBM+22] 通过计算FLOPs来估算计算量。然而,这不适用于主要由整数计算主导的1-bit模型。此外,它主要衡量训练计算,而不是推理计算。为了更好地理解神经语言模型的缩放效率,我们引入了推理优化的缩放法则。它根据能量消耗预测损失。我们关注推理能量成本,因为它随着模型的使用而扩展,而训练成本只有一次。我们按照第2.3节中的方法估算能量消耗。图3展示了在7nm工艺节点上推理能量成本的缩放曲线。结果证明,BitNet具有更高的缩放效率。在固定的计算预算下,BitNet实现了显著更好的损失表现。同时,为了获得与FP16模型相同的性能,BitNet的推理成本更低。
3.3 下游任务的结果
除了损失,我们还关注BitNet在扩展时的能力。相比于损失,由于神经语言模型的涌现性质,能力更难预测。为了用可解释的指标评估能力,我们测试了四个下游任务上的0-shot和4-shot结果,包括Hellaswag [ZHB+19]、Winogrande [SBBC20]、Winograd [LDM12]和Storycloze [MCH+16]。图4报告了不同规模的BitNet和FP16 Transformer的平均结果。类似于损失缩放曲线,下游任务的表现可以随着计算预算的增长而扩展。此外,无论是零样本还是少样本性能,BitNet的能力扩展效率都远高于FP16 Transformer基线。
3.4 稳定性测试
训练低比特Transformer的主要挑战是优化的稳定性。因此,我们通过训练一系列不同峰值学习率的模型,进行了BitNet和FP16基线的稳定性测试。图5a展示了稳定性测试的结果。结果显示,BitNet可以在大学习率下收敛,而FP16 Transformer则不能,展示了BitNet更好的训练稳定性。这种优化优势使得可以使用更大的学习率进行训练。图5b显示,BitNet可以从学习率的增加中受益,在PPL(困惑度)方面实现更好的收敛。
4 与训练后量化的比较
4.1 设置
我们使用与第3.1节中描述的相同设置训练 BitNet。我们将 BitNet 与最新的量化方法进行比较,包括 Absmax [DLBZ22]、SmoothQuant [XLS+23]、GPTQ [FAHA23] 和 QuIP [CCKS23]。这些方法是在 FP16 Transformer 模型上进行训练后量化,遵循与 BitNet 相同的训练设置和数据。其中,Absmax 和 SmoothQuant 量化了权重和激活值,而 GPTQ 和 QuIP 只降低了权重的精度。我们将这些方法应用于各种量化级别。对于仅量化权重的方法(即 GPTQ 和 QuIP),我们尝试了 W4A16 和 W2A16。对于量化权重和激活值的方法(即 Absmax 和 SmoothQuant),我们使用它们将 FP16 Transformers 量化为 W8A8、W4A4 和 W1A8。我们实现的 BitNet 是二进制权重 8 位激活(W1A8),其位数低于或等于基线方法。
4.2 结果
表3详细比较了我们提出的方法 BitNet 与各种基线方法在四个基准数据集(即 Winogrande、Winograd、Storycloze 和 Hellaswag)上的零样本性能。所有模型的大小均为 6.7B,以确保公平比较。这些方法在多个权重位级别上进行了评估,范围从 16 位到 1 位。除了下游任务的零样本准确性外,评估指标还包括验证集上的语言模型困惑度(perplexity),这为每种方法的性能提供了综合理解。
结果表明,BitNet 在实现与基线方法相当的性能水平方面表现出色,特别是在较低位级别时。BitNet 的零样本得分与8位模型相当,而推理成本则低得多。对于4位模型,仅量化权重的方法优于量化权重和激活值的方法,主要原因是激活值更难量化。作为1位模型,BitNet 显著优于量化权重和激活值的方法以及仅量化权重的方法。对于更低位的模型,BitNet 在各个基线方法中表现出一致的优越性。这证明了量化感知训练方法优于训练后量化方法。图6总结了我们的方法与基线方法从1.3B到6.7B模型规模的零样本准确性和少样本准确性。这证明了这种优势在不同规模上是一致的。
5 消融研究
在表4中,我们展示了与几种替代方法的消融研究。我们消融了不同激活量化方法以及模型训练稳定化技术的影响。BitNet 实现了 absmax 来量化激活值,并使用 SubLN 来稳定训练。一种量化替代方法是弹性函数 [LOP+22],它通过可学习参数动态调整比例。在我们的实验中,我们发现 absmax 比弹性函数性能更好。此外,absmax 函数导致更稳定的训练,从而使得 BitNet 可以使用更大的学习率。我们还将 SubLN 与 Pre-LN 和 BMT 架构 [ZGC+23] 进行了比较。Pre-LN 是 GPT 预训练的默认架构,而 BMT 被证明可以提高二值化模型的稳定性。我们的实验表明,SubLN 优于 Pre-LN 和 BMT。因此,我们选择 absmax 和 SubLN 作为 BitNet 的实现。
6 结论与未来工作
我们提出了 BitNet,一种新颖的 1-bit Transformer 架构,用于大型语言模型。我们的方法旨在具有可扩展性和稳定性,能够高效处理大型语言模型。实验结果表明,BitNet 在困惑度和下游任务性能方面实现了具有竞争力的表现,同时显著减少了与基线方法相比的内存占用和能耗。此外,BitNet 遵循了类似于全精度 Transformers 的扩展法则,这表明它可以有效地扩展到更大的语言模型,并在性能和效率方面带来潜在的好处。未来,我们希望在模型规模和训练步骤方面扩大 BitNet 的应用。我们还对将 BitNet 应用于其它架构(例如 RetNet [SDH+23])以训练大型语言模型感兴趣。