BitNet

BitNet 使用专门的 BitLinear 层替换多头注意力机制和前馈网络中的传统线性层。BitLinear 层使用三值精度（取值为 -1、0 和 1）对权重进行量化，并将激活值量化为 8 位精度。

BitNet 模型无法在推理时进行动态量化。由于这是一种量化感知训练（QAT）技术，必须在预训练或微调阶段完成量化。在训练过程中，权重会通过张量级对称量化被量化为三值。

1. 计算权重矩阵绝对值的平均值，并将其用作缩放因子。
2. 将权重除以缩放因子，对数值进行舍入，将其约束在 -1 到 1 之间，然后重新缩放以继续以全精度进行计算。
3. 激活值使用 absmax 量化（对称逐通道量化）被量化到指定的位宽（8 位）。这涉及将激活值缩放到 [−128,127] 范围内。

参考此 PR，使用 Nanotron 对 1.58-bit 模型进行预训练或微调。对于微调，需将 Hugging Face 格式的模型转换为 Nanotron 格式。转换步骤见此 PR。

使用 from_pretrained() 加载 BitNet 量化模型。

from transformers import AutoModelForCausalLM
path = "/path/to/model"
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")

内核

@torch.compile 用于解包权重并执行前向传播。其实现非常直接，并能带来显著的速度提升。未来版本将集成更多优化内核。

资源

阅读 Fine-tuning LLMs to 1.58bit: extreme quantization made easy 以了解 BitNet 模型的训练和微调方式。