BitNet

BitNet 使用专门的 BitLinear 层替换了 Multi-Head Attention 和前馈网络中的传统线性层。BitLinear 层使用三元精度（值为 -1、0 和 1）量化权重，并将激活量化为 8 位精度。

BitNet 模型无法即时量化。它们需要在预训练或微调期间进行量化，因为它是一种量化感知训练 (QAT) 技术。在训练期间，权重使用对称的每张量量化量化为三元值。

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

请参阅此 PR，以使用 Nanotron 预训练或微调 1.58 位模型。对于微调，请将 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 模型是如何训练和微调的。