BitNet b1.58: 革命性的1比特语言模型，性能媲美全精度Transformer

Abstract

1.研究主题：BitNet b1.58 - 新型1比特大型语言模型（LLM）。

2.参数特性：每个参数是三元的（取值为-1, 0, 1）。

3.性能对比：与全精度Transformer LLM（FP16或BF16）模型大小和训练tokens相等，困惑度和端到端任务性能相同。

4.成本效益：在延迟、内存、吞吐量和能耗方面更加高效。

5.新颖贡献：

–定义了LLM的新扩展法则和训练配方。

–提高了性能和成本效益。

–提出了新的计算范式。

–为1比特LLM优化的硬件设计提供了可能性。

1. 1-bit LLMs的时代

1.AI发展：近年来，大型语言模型（LLMs）在规模和能力上迅速增长。

2.部署挑战：由于LLMs规模的增加，部署时对计算能力和环境的考虑。

3.量化技术：采用训练后量化技术创建低比特模型以减少内存和能耗。

4.传统LLMs：通常为16比特浮点数（FP16或BF16），计算成本主要来自浮点加法和乘法。

5.BitNet矩阵乘法：仅涉及整数加法，显著节省能源成本，并转化为更快的计算。

6.参数传输：模型参数从DRAM到芯片内存的传输过程可在推理期间昂贵。

7.1比特LLMs的优势：较低的内存占用和DRAM带宽需求，降低了权重从DRAM加载的成本和时间。

8.BitNet b1.58介绍：每个参数取值在{-1, 0, 1}中，增加了0的取值，没有乘法操作需求，大幅提升内存消耗和能耗效率。

9.性能：BitNet b1.58在困惑度和端到端任务性能上与FP16 LLM基线相当，即使是在3B参数的配置下。

图-1展示了1比特大型语言模型（例如BitNet b1.58）与传统的16比特浮点数（FP16/BF16）Transformer LLMs在性能和成本之间提供的帕累托改进（Pareto Improvement）。 图中的上半部分突出显示了BitNet b1.58在保持与传统Transformer LLMs相同性能的同时，如何降低推理成本（包括延迟、吞吐量和能量）。它指出BitNet b1.58的新计算范式需要为1比特LLMs设计新的优化硬件。 图中的下半部分给出了一个计算示例，比较了FP16模型和1比特模型在处理矩阵乘法时的差异： - 在FP16模型中，权重矩阵（W）包含16比特的浮点数，输入（X）与权重矩阵相乘，产生输出（Y）。这一过程涉及到大量的乘法和加法运算，需要在GPU上执行。 - 在1比特模型中，权重矩阵只包含三个可能的值{-1, 0, 1}，这大大简化了乘法运算，因为这些乘法实际上变成了带符号的加法或减法（或者在乘以0时不做任何操作）。这减少了对GPU的依赖，并为新硬件的开发铺平了道路，这些硬件将专门为1比特计算进行优化。 这张图强调了1比特模型如何通过减少计算复杂度来提高效率，同时为未来在能效和成本方面更为高效的LLM硬件设计提供了动力。

2. BitNet b1.58

2.1 BitNet b1.58概述

•架构：基于BitNet架构，是一种替换了nn.Linear为BitLinear的Transformer。

•训练：从零开始训练，采用1.58比特权重和8比特激活函数。

•改进：相比原始BitNet，引入了一些修改。

2.2 量化函数

•目的：将权重约束在-1, 0, 或 +1。

•方法：采用absmean量化函数，通过权重矩阵的平均绝对值来缩放权重矩阵，然后将每个值四舍五入到最近的整数{-1, 0, +1}。

•公式：

–

，使用RoundClip函数来确保权重值在允许的范围内。

–

–

，缩放因子

，它是权重矩阵所有元素绝对值的平均值。

2.3 激活函数量化

•实现：与BitNet相同，但不在非线性函数前对激活函数进行缩放。

•调整：所有激活函数都缩放到

，以消除零点量化。

>>> 举例讲解：量化函数 与 激活函数量化

假设我们有一个简化的权重矩阵

和一个激活函数的输出向量

，并且我们想要应用BitNet b1.58的量化方法。

步骤1：定义权重矩阵和激活函数输出

假设

是一个

权重矩阵:

激活函数的输出

是一个向量:

步骤2：计算权重矩阵的平均绝对值

计算

值:

步骤3：应用量化函数

使用公式（1）对

进行量化。我们需要在量化之前添加一个小的常数

来避免除以0的情况，这里我们假设

。

现在，我们计算

:

对于

中的每一个元素，我们需要除以

，并四舍五入到最近的整数 {-1, 0, +1}。

例如:

对于矩阵

中的所有元素重复这个过程，我们得到量化后的权重矩阵

。

步骤4：量化激活函数输出

根据文档，激活函数输出的量化不需要像权重那样缩放，但是我们需要将它们缩放到特定的范围。假设我们选择的范围是 ([-2, 2])（即

），我们将每个激活值缩放到这个范围，并四舍五入到最近的整数。

例如:

所以量化后的激活向量

是:

这样，我们就完成了权重矩阵和激活函数输出的量化过程。这两个量化过程减少了模型的存储需求，并且可能提高计算效率，因为它们减少了乘法运算的需要，并允许使用整数计算。

2.4 LLaMA-类组件

•背景：LLaMA架构已成为开源LLMs的事实上的骨干。

•集成：BitNet b1.58采用LLaMA-类组件，如RMSNorm和rotary embedding，并移除了所有偏差。

•兼容性：BitNet b1.58可以轻松集成到流行的开源软件中，如Huggingface和vLLM。

3 Results

3.1 实验比较

•比较对象：BitNet b1.58与重构的FP16 LLaMA LLM。

•数据集：在RedPajama数据集上预训练了1000亿tokens。

•性能评估：在多个语言任务上进行零样本(zero-shot)性能评估，包括ARC-Challenge, Hellaswag, Winogrande, PIQA, OpenBookQA, BoolQ, WikiText2和C4数据集。

•性能测量：使用FasterTransformer代码库测量GPU内存、延迟和每输出token的时间。

3.2 实验结果

•BitNet b1.58在3B模型大小时，开始匹配LLaMA LLM的全精度性能，同时速度提高2.71倍，使用的GPU内存减少3.55倍。

•在3.9B模型大小时，BitNet b1.58在性能上显著优于LLaMA LLM 3B，内存和延迟成本更低。

3.3 扩展模型大小实验

•评估了7B, 13B, 和70B大小的模型。

•发现随着模型大小增加，BitNet b1.58的速度优势增加。

•在70B模型上，BitNet b1.58比LLaMA LLM基线快4.1倍，内存使用也减少了相应比例。

•能源消耗：BitNet b1.58在矩阵乘法操作上比LLaMA LLM节省71.4倍的能源。

•吞吐量：BitNet b1.58 70B可以支持比LLaMA LLM的批处理大小大11倍，实现更高的吞吐量。

3.4 新的扩展规律

•13B BitNet b1.58比3B FP16 LLM在延迟、内存使用和能耗方面更高效。

•30B BitNet b1.58比7B FP16 LLM更高效。

•70B BitNet b1.58比13B FP16 LLM更高效。

3.5 2T tokens训练

•BitNet b1.58在2T tokens上的扩展能力测试。

•在包括Winogrande和ARC-easy在内的基准测试中，BitNet b1.58的性能优于StableLM 3b模型。

4 讨论与未来发展方向

1比特专家混合体（MoE）LLMs

•MoE已被证明是LLMs的一种成本效益高的方法。

•MoE的挑战包括高计算FLOPs、高内存消耗和芯片间通信开销。

•1.58比特LLMs可以减少部署MoE模型所需的设备数量，并显著减少网络传输激活的开销。

LLMs对长序列的原生支持

•处理长序列能力是LLMs的关键需求。

•BitNet b1.58通过将激活从16比特减少到8比特，使得在相同资源下上下文长度翻倍。

•这可以进一步无损压缩到4比特或更低，为1.58比特LLMs留下了未来工作的空间。

边缘和移动设备上的LLMs

•1.58比特LLMs有可能大幅改善边缘和移动设备上的语言模型性能。

•这些设备通常受到内存和计算能力的限制。

•降低内存和能耗使得1.58比特LLMs能够部署在这些设备上，从而开启了以前不可能的应用范围。

1比特LLMs的新硬件

•最近的工作（如Groq）展示了为LLMs构建特定硬件（如LPUs）的有希望的结果和潜力。

•鉴于新的计算范式，我们设想并呼吁设计专门为1比特LLMs优化的新硬件和系统。