本文涉及到的详细测试代码和测试步骤放置于:
https://github.com/xinyuwei-david/david-share.git下的:DeepLearning/1-Bit- LLMs
本文中不再赘述代码实现。欢迎给repo点亮Star,您的点赞是作者持续创作的动力。
先放个实测的视频,在笔记本上实现纯用CPU推理Llama3-8B。后面如果有专门支持三元权重的CPU/NPU芯片出来,BitNet应该会有用武之地。
在当前的人工智能领域,大型语言模型(LLM)已经成为了不可或缺的组成部分。然而,这些模型庞大的参数数量带来了巨大的计算和内存需求,限制了其在资源有限设备上的应用。那么,是否有可能在不显著降低模型性能的情况下,减少模型的计算和存储需求呢?微软研究院的 BitNet 架构和 1-Bit LLMs 给出了肯定的答案。
本文将结合多个角度,深入探讨 1-Bit LLMs 的概念及其实现方式,并通过丰富的例子帮助读者更好地理解这一前沿技术。
一、什么是 1-Bit LLMs?
1-Bit LLMs,直观上理解,就是使用低位宽(位数)的模型参数(权重)来表示和计算的大型语言模型。具体来说,1-Bit LLMs 中的每个权重被限制为三元值:-1、0、1。这种表示方法显著降低了模型的内存占用和计算复杂度。
虽然被称为“1-Bit”,但实际上每个权重的平均占用位数约为 1.58 位。这个值是通过信息论中的熵计算得出的,表示了在理论上表示三元值所需的最小平均位数。
1-Bit LLMs 与 FP16/FP32 的区别
在深度学习中,我们经常提到 FP16 和 FP32,这是指模型中使用的浮点数数据类型的精度。其中:
FP32(32 位浮点数):每个浮点数占用 32 位,具有高精度和宽动态范围。
FP16(16 位浮点数):每个浮点数占用 16 位,占用内存更小,计算速度更快,但精度较低。
1-Bit LLMs 与之不同,主要强调的是模型权重的位宽缩减。通过将权重量化为三元值(-1、0、1),并使用低位宽的数据类型(平均每个权重约占 1.58 位)进行存储和计算,达到了更极致的内存和计算效率。
在传统的 FP16 或 FP32 模型中,模型参数(权重和偏置)通常以 16 位或 32 位浮点数形式存储和计算。而在 1-Bit LLMs 中,模型参数被量化为三元值,在运行时以低位宽的数据类型表示,大大降低了内存占用和计算复杂度。
1-Bit LLMs 主要是将模型的权重参数量化为接近 1 位的表示形式(平均每个权重约占 1.58 位),即使用三元值(-1、0、1)来表示权重。这一做法显著降低了模型在运行时对内存和计算的需求。
然而,需要注意的是,模型的其他部分仍然可能使用更高精度的数据类型。具体来说:
激活值(Activations):
表示方式:在 1-Bit LLMs 中,激活值通常被量化为 8 位 或者 16 位。这确保了模型在前向传播过程中保持足够的数值精度,防止信息在层与层之间的传递中损失过多。
原因:激活值过度量化可能导致模型性能显著下降,因为激活值在网络的各层之间传递,是模型表达能力的关键部分。
梯度(Gradients):
表示方式:在训练过程中,梯度计算通常需要更高的精度,可能使用 16 位(FP16) 或 32 位(FP32) 浮点数。
原因:梯度的精度对于模型的训练稳定性至关重要。过低的精度可能导致梯度下溢或上溢,影响模型的收敛性和最终性能。
优化器状态(Optimizer States):
表示方式:优化器内部维护的状态(如动量项、一阶和二阶动量等)通常也使用 FP32 进行存储。
原因:这些状态涉及到参数更新的细微调整,需要高精度来准确捕捉参数变化。
损失值和其他计算:
表示方式:损失值、正则化项等关键计算通常保持在 FP32 精度,以确保训练过程的数值稳定性。
因此,1-Bit LLMs 主要是在模型的权重部分实现了位宽的极致压缩,而模型的其他部分 (如激活值、梯度、优化器状态等) 仍然使用更高的位宽来保持模型的性能和训练的稳定性。举个例子:
假设我们有一个传统的神经网络模型:权重(Weights):使用 FP32 表示,每个权重占用 32 位。
激活值(Activations):使用 FP32。
梯度(Gradients):使用 FP32。
在 1-Bit LLMs 中:权重:量化为三元值 (-1, 0, 1),平均每个权重占用约 1.58 位。
激活值:可能量化为 8 位整数(INT8)。
梯度:可能仍然使用 FP16 或 FP32。
这样做的好处是:大幅减少了模型的存储需求和计算复杂度,因为权重参数通常占据了模型大部分的存储空间。
保持了模型的性能和训练稳定性,因为关键的计算部分仍然使用较高的精度。
需要注意的点:模型的训练过程可能更加复杂:由于权重被量化,训练时需要使用特定的技术(如直通估计器 STE)来解决量化带来的不可微问题。
硬件支持:现有的硬件可能需要特定的优化才能充分利用这种低位宽的权重表示。
结论:1-Bit LLMs 并不是将整个模型都转换为 1 位表示,而是主要针对模型的权重部分。
其他部分仍然保留较高的位宽,以确保模型的正常功能和性能。
二、为什么使用三元权重?
在神经网络模型,尤其是大型语言模型(LLM)中,权重在推理过程中会占用大量的内存。甚至是占用推理中的绝大部分内存。
原因如下:
模型规模庞大:大型语言模型通常包含数亿、数十亿甚至数千亿个参数(权重和偏置)。这些参数需要在推理时加载到内存中,以进行前向传播计算。
数据类型的精度:权重通常以浮点数形式存储,如 32 位浮点数(FP32)或 16 位浮点数(FP16)。高精度的数据类型意味着每个权重占用更多的内存空间。
计算需求:在推理过程中,模型需要访问和计算所有相关的权重,以生成输出结果。这进一步增加了内存和计算资源的消耗。
传统的神经网络模型通常使用 32 位(FP32)或 16 位(FP16)的浮点数来表示权重和激活值。这些高精度的表示方式虽然能够提供丰富的表达能力,但也带来了巨大的存储和计算需求。
使用三元权重的主要优势包括:
内存效率:每个权重只需约 1.58 位,比 FP32 的 32 位和 FP16 的 16 位大大减少了存储需求。
计算效率:三元权重的乘法运算可以大幅简化,乘以 0 的操作可以跳过,乘以 1 或 -1 的操作可以简化为复制或取反。
模型小型化:使大型模型可以在普通 CPU 或资源受限的设备上运行,扩大了模型的应用范围。
三、BitNet 架构:实现 1-Bit LLMs 的关键
微软研究院开发的 BitNet 架构是实现 1-Bit LLMs 的核心。它通过以下关键技术实现了模型的高效性:
1. 三元权重量化
在 BitNet 中,模型的权重被量化为三元值(-1、0、1)。具体的量化过程如下:
计算缩放因子:首先,计算权重矩阵的平均绝对值,作为缩放因子 gamma。
计算公式如下:
gamma = (1 / (n * m)) * 所有权重绝对值之和
其中,n 和 m 分别是权重矩阵的行数和列数。
缩放权重矩阵:将权重矩阵中的每个元素除以缩放因子 gamma,得到缩放后的权重矩阵。
应用 RoundClip 函数:对缩放后的权重矩阵中的每个元素,使用 RoundClip 函数进行量化:
RoundClip(x, -1, 1) = max(-1, min(1, round(x)))
该函数将值四舍五入到最接近的整数,并限制在 -1 到 1 之间。
示例:
假设有一个权重矩阵:
W = | 0.3 -0.7 |
| 0.5 0.1 |
计算缩放因子 gamma:
gamma = (1 / (2 * 2)) * ( |0.3| + |-0.7| + |0.5| + |0.1| )
= 0.4
缩放权重矩阵:
缩放后的权重矩阵 = W / gamma
= | 0.3 / 0.4 -0.7 / 0.4 |
| 0.5 / 0.4 0.1 / 0.4 |
= | 0.75 -1.75 |
| 1.25 0.25 |
应用 RoundClip 函数:
得到量化后的权重矩阵:
W_quantized = | 1 -1 |
| 1 0 |对于 0.75:
RoundClip(0.75, -1, 1) = 1
对于 -1.75:
RoundClip(-1.75, -1, 1) = -1
对于 1.25:
RoundClip(1.25, -1, 1) = 1
对于 0.25:
RoundClip(0.25, -1, 1) = 0
2. 直通估计器(STE)
由于量化函数的非连续性和不可微性,标准的反向传播无法直接应用。为了解决这个问题,BitNet 使用了 直通估计器(Straight Through Estimator, STE)。
原理:在前向传播中,使用离散的三元权重进行计算;在反向传播中,假设量化函数是可微的,将梯度直接传递给未量化的连续权重。
效果:使得模型能够在训练过程中更新权重,同时保持权重的离散性。
3. 激活值量化
除了权重量化外,为了进一步提高模型的效率,BitNet 还将激活值量化为 8 位精度。
层归一化:在量化激活值之前,应用层归一化(Layer Normalization),确保输出的稳定性,防止梯度爆炸或消失。
缩放和量化:计算激活值的最大绝对值,作为缩放因子,将激活值缩放到适合 8 位数表示的范围,然后进行量化。
四、优化的推理内核:bitnet.cpp
由于现有硬件并未针对三元权重优化,微软开发了 bitnet.cpp,包含了多个优化的推理内核,如 I2_S、TL1 和 TL2。
I2_S 内核:将三元权重打包为 2 位表示,适用于多线程的 CPU 环境。
TL1 内核:将每两个权重打包,使用查找表加速计算,适用于线程数有限的环境。
TL2 内核:进一步压缩权重,适用于内存和带宽受限的环境。
通过这些优化,bitnet.cpp 实现了在标准 CPU 上高效地运行 1-Bit LLMs。
五、神经网络模型的本质:静态部分与运行时状态
要深入理解 1-Bit LLMs,我们需要回顾神经网络模型的本质,可以从 静态部分 和 运行时状态 两个角度来看。
1. 静态部分
模型权重文件(如
model_weights.h5
或model_weights.pt
):存储了模型的权重和偏置。这些参数在训练过程中通过优化算法学习得到。
配置文件(如
config.json
):定义了模型的架构,包括层数、每层的类型、层的参数等。
词汇表文件(如
vocab.txt
):在处理文本数据时使用,存储了词汇和它们的索引,用于将文本转换为模型可以处理的数值形式。
附加工具和脚本(如
preprocess.py
、postprocess.py
):用于数据的预处理和后处理,如图像的大小调整、归一化,或文本的分词等。
训练/微调脚本(如
train.py
、finetune.py
):包含用于训练或微调模型的代码,这些脚本定义了训练过程,包括损失函数的选择、优化器的配置和训练周期的设置。
2. 运行时部分
输入数据(Input):
模型接收的原始数据,可以是各种形式,如图片、文本或音频等。
激活值(Activations):
每一层通过激活函数处理后的输出值,作为下一层的输入。
中间状态(Intermediate States):
每一层可能产生的其他形式的输出,如卷积层中的特征图。
梯度(Gradients):
在训练过程中,通过反向传播算法计算得到的每个参数的梯度。
损失值(Loss Value):
损失函数计算得到的当前模型输出与真实标签之间的差异值。
状态更新(State Updates):
根据梯度和学习率更新的权重,以及优化器维护的状态(如动量和自适应学习率信息)。
缓存(Caches):
在前向传播过程中缓存的某些值,用于在反向传播中有效地计算梯度。
通过理解神经网络模型的静态结构和运行时状态,我们可以更好地理解 1-Bit LLMs 是如何在保持模型性能的前提下,通过优化权重和激活值的位宽,达到高效运行的目的。
3. 模型参数与运行时的关系
模型参数通常指的是神经网络模型中的权重(weights)和偏置(biases),这些参数是在训练过程中学习得到的,用于决定模型的行为和性能。
模型参数既存在于静态存储中,也存在于运行时内存中。
1)静态部分的模型参数
存储在磁盘上的文件:
权重文件:模型训练完成后,权重和偏置参数会被保存到磁盘上的文件中,例如
.pt
、.h5
、.ckpt
等格式。存储格式:这些参数可以以不同的数据类型和精度存储,例如 FP32、FP16、INT8 等。
静态模型参数的作用:
持久化存储:保存模型的当前状态,便于后续的加载、部署或分享。
无需在磁盘上进行计算:存储的参数本身只是数据,没有参与计算。
2)运行时的模型参数
加载到内存中:
当您加载模型时(例如,使用
model.load_state_dict()
),模型参数被从磁盘读取并加载到计算设备的内存中(如 RAM、GPU 显存)。运行时的表示和计算:
推理阶段:模型参数与输入数据进行计算,生成输出结果。
训练阶段:模型参数在计算梯度后根据优化算法进行更新。
FP32:32 位浮点数,传统上用于高精度计算。
FP16:16 位浮点数,减少内存占用和提高计算速度。
低位宽格式:如 8 位整数(INT8)、1 位二值(Binary)或三值(Ternary)表示,用于进一步压缩模型并加速计算。
数据类型:模型参数在内存中以特定的数据类型表示,可能与磁盘上的存储格式相同或不同。
参与计算:加载到内存后的模型参数将被用于前向传播和反向传播计算。
运行时的数据类型影响:
示例:使用 FP16 或 INT8 可以减少内存带宽,占用更少的缓存,提高处理器的并行计算能力。
低精度可能导致精度下降,但通过适当的训练和量化方法,可以减小影响。
计算精度:数据类型的精度会影响计算结果的准确性。
计算性能:低位宽的数据类型通常可以提高计算速度和效率。
3)模型参数与运行时的关系
当我们讨论模型参数时,既包括它们在磁盘上的存储,也包括它们在内存中的运行时表示。
重点在于运行时的使用:
计算中的作用:模型参数在运行时被加载到内存中,参与实际的计算过程,这时的数据类型和表示方式对模型的性能和精度至关重要。
优化和量化:为了提高运行时的效率,我们可以对模型参数进行量化,使用更低的位宽表示,例如 1-Bit LLMs 中的三元权重(-1、0、1)。
4)1-Bit LLMs 中的模型参数
权重量化:
三元权重表示:在 1-Bit LLMs 中,模型参数被量化为三元值(-1、0、1),平均每个权重占用约 1.58 位。
运行时计算:在推理或训练过程中,这些量化后的权重在内存中以低位宽的数据类型表示,并用于实际的计算。
与运行时的关系:
模型性能:量化后的参数在运行时直接影响模型的计算性能和占用的内存资源。
硬件支持:为有效利用这些低位宽参数,可能需要特定的硬件支持或优化的软件实现,以加速计算。
5)总结
模型参数既指存储在磁盘上的静态参数,也指加载到内存后在运行时参与计算的参数。
当讨论模型参数的数据类型或位宽(如 1-Bit LLMs 中的三元权重)时,通常是指参数在运行时**加载到内存后的表示方式和计算方式。
通过理解模型参数与运行时的关系,我们可以更好地理解 1-Bit LLMs 的优势所在。模型参数的位宽和数据类型直接影响模型的计算性能、内存占用和精度。1-Bit LLMs 通过极大地降低模型参数的位宽,实现了在运行时的高效计算和低内存占用。
六、总结与展望
1-Bit LLMs 通过创新的权重量化和计算优化方法,显著降低了模型的存储和计算需求,使得在资源受限的环境中运行大型语言模型成为可能。
然而,使用三元权重也带来了一定的挑战:
表达能力受限:权重取值的限制可能影响模型的表达能力,需要通过增加网络深度或改进训练方法来弥补。
硬件支持:现有硬件尚未针对三元权重进行优化,需要专门的软件和内核支持。
训练复杂性:需要特殊的训练方法,如 STE,来处理量化带来的非连续性。
尽管如此,1-Bit LLMs 的出现为模型压缩和高效计算开辟了新的道路。随着技术的进一步发展和硬件支持的提升,我们有理由相信,这种高效的模型将在更多的实际应用中发挥重要作用。