BitNet
BitNet 在多头注意力机制和前馈网络中,用称为 BitLinear 的专用层替换了传统的线性层,这些层使用三元(或旧版本中的二元)精度。这里引入的 BitLinear 层使用三元精度(值为 -1、0 和 1)对权重进行量化,并将激活量化为 8 位精度。
![Alt Text](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/1.58llm_extreme_quantization/bitlinear.png)
在训练过程中,我们首先使用对称的每张量量化将权重量化为三元值。首先,我们计算权重矩阵绝对值的平均值,并将其用作比例。然后,我们将权重除以比例,对值进行四舍五入,将其限制在-1和1之间,最后重新缩放它们以继续以全精度进行。
激活值随后使用absmax量化(对称每通道量化)量化为指定的位宽(例如8位)。这涉及将激活值缩放到范围[−128,127[。量化公式为:
要了解更多关于我们如何训练和微调bitnet模型的信息,请查看博客文章这里
从Hub加载BitNet模型
BitNet模型不能即时量化——它们需要在应用量化的情况下进行预训练或微调(这是一种量化感知训练技术)。一旦训练完成,这些模型就已经量化,并在中心以打包版本的形式提供。
可以加载量化模型:
from transformers import AutoModelForCausalLM
path = "/path/to/model"
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")
预训练 / 微调 BitNet 模型
如果你想使用Nanotron预训练或微调自己的1.58位模型,请查看这个PR,所有你需要开始的内容都在那里!
为了进行微调,您需要将模型从Hugging Face格式转换为Nanotron格式(两者之间存在一些差异)。您可以在此PR中找到转换步骤。
内核
在我们的初始版本中,我们选择使用@torch.compile
来解包权重并执行前向传递。它的实现非常直接,并且带来了显著的加速改进。我们计划在未来的版本中集成更多优化的内核。