1-bit 大型语言模型（LLM）时代的到来

人工智能咨询培训老师叶梓 转载标明出处

模型规模的扩大带来了部署上的挑战，并因其高能耗引对环境和经济产生了影响。为了应对这些挑战，研究者们开始探索使用低位宽量化技术来降低模型的推理成本，同时保持模型性能。微软公司和中国科学院大学的研究团队提出了一种名为BitNet b1.58的1-bit LLM变体。BitNet b1.58的每个参数（或权重）都是三元的{-1, 0, 1}，与相同模型大小和训练token的全精度（例如FP16或BF16）Transformer LLM相比，在困惑度和最终任务性能方面表现相当，但在延迟、内存、吞吐量和能耗方面具有显著的成本效益。

方法

BitNet b1.58是基于BitNet架构的，后者是一种Transformer模型，其特点是将标准的nn.Linear层替换为BitLinear层。这种替换使得模型能够从头开始训练，使用1.58位权重和8位激活值。与原始的BitNet相比，b1.58版本引入了一些关键的改进。

量化函数是这些改进中的一个关键点。为了将权重限制在-1、0或+1的范围内，研究者们采用了平均绝对值量化函数（absmean quantization function）。这个函数首先通过权重矩阵的平均绝对值进行缩放，然后对每个值四舍五入到最近的整数，即{-1, 0, +1}集合中的一个。量化函数的数学表达式如下：

这里，RoundClip函数确保了量化值在-1和1之间，并且四舍五入到最接近的整数。γ是权重矩阵平均绝对值的缩放因子，而ε是用于防止四舍五入误差的小量。

与原始的BitNet不同，BitNet b1.58在激活函数的量化上没有将激活值缩放到[0, Qb]的范围，而是将每个token的激活值缩放到[−Qb, Qb]。这种处理方式简化了实现和系统级优化，同时实验中对性能的影响微乎其微。

为了更好地融入开源社区，BitNet b1.58的设计采用了与LLaMA相似的组件。它使用了RMSNorm、SwiGLU、旋转嵌入等技术，并去除了所有偏置项。这使BitNet b1.58可以轻松地集成到流行的开源软件中，如Huggingface、vLLM和llama.cpp等。

BitNet b1.58的性能评估与比较

研究者们对BitNet b1.58与FP16精度的LLaMA LLM进行了多尺寸的比较。为了保证比较的公平性，所有模型都在RedPajama数据集上进行了预训练，处理了高达1000亿个token。研究者们评估了模型在多种语言任务上的零样本性能，包括但不限于ARC-Easy、ARC-Challenge、Hellaswag、WinoGrande、PIQA、OpenbookQA和BoolQ等。另外还在WikiText2和C4数据集上报告了模型的验证困惑度。

在评估过程中，研究者们使用了FasterTransformer代码库来测量LLM在GPU设备上的推理延迟，同时为BitNet b1.58集成了Ladder的2位内核。重点报告了每个输出token的时间，因为这代表了推理的主要成本。

Table 1汇总了BitNet b1.58和LLaMA LLM的困惑度和成本。结果显示，当模型大小达到3B时，BitNet b1.58在困惑度方面开始与全精度的LLaMA LLM相匹配，同时在速度上快了2.71倍，GPU内存的使用量减少了3.55倍。特别是，3.9B大小的BitNet b1.58在速度上快了2.4倍，内存消耗少了3.32倍，但性能明显优于3B的LLaMA LLM。

Table 2展示了BitNet b1.58和LLaMA LLM在不同任务上的零样本学习能力。这些任务包括ARCe、ARCc、HS、BQ、OQ、PQ、WGe等。可以看出，随着模型大小的增加，BitNet b1.58的性能逐渐接近甚至在某些情况下超过了LLaMA LLM。特别是在3.9B模型大小上，BitNet b1.58在所有任务上都显示出了更高的平均准确率。

研究者进一步扩大了模型尺寸至7B、13B和70B，并评估了成本。Figure 2展示了随着模型尺寸的增加，延迟和内存消耗的趋势，显示出随着模型尺寸的增加，加速比也在增加。特别是，70B的BitNet b1.58比LLaMA LLM基线快了4.1倍。这是因为nn.Linear的时间成本随着模型尺寸的增加而增加。内存消耗也呈现类似的趋势，因为嵌入保持全精度，其内存比例对于更大的模型来说更小。延迟和内存的测量都是使用2位内核完成的，因此还有进一步降低成本的优化空间。

研究者们还估算了BitNet b1.58和LLaMA LLM的算术运算能耗，主要集中在矩阵乘法的计算上，因为这是LLMs成本的主要部分。Figure 3展示了能耗成本的组成。BitNet b1.58的大部分是INT8加法计算，而LLaMA LLM包括FP16加法和FP16乘法。根据[Hor14, ZZL22]中的能耗模型，BitNet b1.58在7nm芯片上节省了71.4倍的矩阵乘法算术运算能耗。研究者们还报告了具有512个token的模型的端到端能耗成本。结果显示，随着模型尺寸的增加，BitNet b1.58在能耗方面比FP16 LLaMA LLM基线更加高效。

研究者们比较了具有70B参数的BitNet b1.58和LLaMA LLM在两个80GB A100卡上的吞吐量，使用流水线并行[HCB+19]，以便LLaMA LLM 70B能够在设备上运行。他们增加了批量大小，直到达到GPU内存限制，并保持序列长度为512。Table 3显示，BitNet b1.58 70B可以支持LLaMA LLM批量大小的11倍，从而实现了8.9倍的更高吞吐量。

为了测试BitNet b1.58在token数量上的可扩展性，研究者们训练了一个使用2T token的BitNet b1.58模型，遵循StableLM-3B的数据配方，这是最先进的开源3B模型。两个模型都在由WinoGrande、PIQA、SciQ、LAMBADA和ARC-easy组成的基准上进行了评估。Table 4报告了零样本准确率。对于使用准确率和归一化准确率测量的任务，研究者们取了两者的平均值。StableLM 3b在2T token的结果直接取自其技术报告。研究结果表明，BitNet b1.58在所有最终任务上都取得了更优越的性能，表明1.58位LLMs也具有强大的泛化能力。

这些发现表明，BitNet b1.58是在现有最先进LLM模型基础上的帕累托改进，为大模型的能效优化和性能提升开辟了新的道路。

论文链接：https://arxiv.org/abs/2402.17764

项目链接：Advancing AI for humanity | Foundation of AI