Transformers 文档

BitNet

BitNet

BitNet 在多头注意力机制和前馈网络中,用称为 BitLinear 的专用层替换了传统的线性层,这些层使用三元(或旧版本中的二元)精度。这里引入的 BitLinear 层使用三元精度(值为 -1、0 和 1)对权重进行量化,并将激活量化为 8 位精度。

Alt Text
The architecture of BitNet with BitLinear layers

在训练过程中,我们首先使用对称的每张量量化将权重量化为三元值。首先,我们计算权重矩阵绝对值的平均值,并将其用作比例。然后,我们将权重除以比例,对值进行四舍五入,将其限制在-1和1之间,最后重新缩放它们以继续以全精度进行。 scalew=11nmijWij scale_w = \frac{1}{\frac{1}{nm} \sum_{ij} |W_{ij}|} Wq=clamp[1,1](round(Wscale)) W_q = \text{clamp}_{[-1,1]}(\text{round}(W*scale)) Wdequantized=Wqscalew W_{dequantized} = W_q*scale_w

激活值随后使用absmax量化(对称每通道量化)量化为指定的位宽(例如8位)。这涉及将激活值缩放到范围[−128,127[。量化公式为: scalex=127Xmax,dim=1 scale_x = \frac{127}{|X|_{\text{max}, \, \text{dim}=-1}} Xq=clamp[128,127](round(Xscale)) X_q = \text{clamp}_{[-128,127]}(\text{round}(X*scale)) Xdequantized=Xqscalex X_{dequantized} = X_q * scale_x

要了解更多关于我们如何训练和微调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来解包权重并执行前向传递。它的实现非常直接,并且带来了显著的加速改进。我们计划在未来的版本中集成更多优化的内核。

< > Update on GitHub