微调 LLM 至 1.58 位：使用 BitNet 的极端量化

Hugging Face 已经开发出一种方法，能够将已有的大型语言模型（LLM）微调至 1.58 位精度，采用 BitNet 架构。该方法只使用三种取值（-1、0、1）来表示参数，极大降低了计算和能耗，而无需像从头预训练 1 位模型那样庞大的预算。

BitNet 架构与 1.58 位量化

BitNet 用 BitLinear 层取代多头注意力和前馈网络中的标准 Linear 层。这些层对权重使用三值精度，对激活使用 8 位精度。

与依赖 FP16 加法和乘法的标准 LLM（例如 Llama）不同，BitNet b1.58 在矩阵乘法中使用 INT8 加法。理论上，这种计算方式相较于 Llama 基线可将矩阵乘法的能耗降低 71.4 倍。

由于用于三值量化的 round() 函数不可导，BitNet 采用 Straight Through Estimator (STE)。STE 将取整操作的梯度近似为 1，使梯度能够像通过恒等函数一样流经该操作，从而支持标准的基于梯度的优化。

权重： 采用对称的 per‑tensor 量化。尺度为权重矩阵绝对值均值的倒数。权重先被缩放、取整、限制在 -1 到 1 之间，然后再重新缩放。
激活： 使用 per‑token 的 absmax 量化到 8 位，将数值缩放至 [-128, 127] 范围。激活量化前先进行层归一化（Layer Normalization，LN），以保持输出方差。

Hugging Face 成功将 Llama 3 8B 模型微调至 1.58 位精度。最初的实验表明，直接引入 BitLinear 层会导致模型几乎失去所有预训练信息，出现损失急剧上升的情况。

为防止先前知识的丢失，Hugging Face 实现了一个动态的 $\lambda$（lambda）值，以逐步引入量化：

$$\lambda = \min\left(\frac{\text{training_step}}{\text{total_training_steps}}, 1\right)$$

通过将原始值与量化后值的差乘以 $\lambda$，模型从全精度（$\lambda=0$）平滑过渡到完全量化（$\lambda=1$）。该线性调度器使收敛更好，在 TinyStories 数据集上实现约 4 的困惑度。

为了确保模型保留通用知识且不在小数据集上过拟合，团队将训练规模扩大到 FineWeb-edu 数据集。使用 1e-4 的学习率、每批 200 万 token，训练 100 亿 token，模型在 WikiText 上达到了 12.2 的困惑度。

进一步扩展到 1000 亿 token 表明，虽然在某些指标上模型表现接近原始 Llama 3 8B，但整体仍略逊于全精度基线。

使用 1.58 位架构微调的模型已在 HF1BitLLM 组织下发布。

竞争性表现： 在微调 100 亿 token 后，1.58 位 Llama 3 8B 模型的表现超过了在 1000 亿 token 上训练的 BitNet 7B 模型以及在 1.26 万亿 token 上蒸馏的 FBI LLM。
MMLU 基准： 所开发的 8B 模型在 MMLU 基准上超越了 Llama 1 7B 模型。
模型体积： 将权重打包为 int8 张量后，参数量从 8B 降至 2.8B。

为实现 1.58 位权重的速度和内存优势，Hugging Face 实现了自定义 CUDA 与 Triton 内核，在矩阵乘法期间进行即时权重解包。

为克服内存带宽瓶颈和冗余数据访问，团队采用 tiling 技术。该技术将矩阵划分为适配 GPU 快速共享内存的更小子矩阵（瓦片），从而降低慢速全局内存访问的频率。

Triton vs. Torch： 自定义 Triton 内核的性能大致等同于使用 BF16 精度的 @torch.compile。
BitBlas： 团队发现 BitBlas（一种混合精度软件库）在低精度下的表现优于自定义 Triton 内核和 Torch 的 matmul，但由于内核编译导致加载时间更长。

集成通过在 transformers 库中新增的 “bitnet” 量化方法实现。标准 Linear 层被专用的 BitLinear 层替换。API 保持不变，用户仍可使用 AutoModelForCausalLM.from_pretrained 加载模型。