BitNet

PyTorch中对论文《BitNet：为大型语言模型扩展1位Transformer》中的线性方法和模型的实现

论文链接：

BitLinear = 张量 -> 层归一化 -> 二值化 -> 绝对最大量化 -> 反量化

“BitNet架构的实现非常简单，只需替换Transformer中的线性投影（即PyTorch中的nn.Linear）即可。”——BitNet真的很容易实现，只需将线性层替换成BitLinear模块！

新闻

• 新迭代 🔥 论文《1位LLM时代：所有大型语言模型都在1.58位内》提出了全新的迭代版本，我们正在实现它。加入Agora Discord并贡献你的力量吧！点击加入
• 新优化 第一个BitLinear已经优化完毕，我们现在有了一个将BitLinear融入注意力机制的BitAttention BitMGQA。多组查询注意力因其快速解码和长上下文处理能力而广受认可，感谢Frank提供的易于使用的实现！
• BitLinear 1.5发布 🔥：新的BitLinear 1.5仍在开发中 🔥 文件在此 目前仍存在一些问题，比如反量化算法，而且我们还需要用逐元素相加替代乘法运算。如果你能在这方面提供帮助，那就太棒了。
• 注意：要使用BitLinear，模型显然需要从头开始进行微调；仅仅在已训练好的模型中更换线性方法是行不通的。请务必进行微调或从零开始训练。

致谢

• Dimitry、Nullonix：感谢他们的分析以及代码审查与修订工作
• Vyom：感谢他提供了4080显卡用于训练！

安装

pip3 install bitnet

使用

我们在本项目及examples文件夹中提供了丰富的示例脚本。如果你在Discord上需要某个用例的帮助，请随时告知我们！

BitLinear

• 这是论文中的主要创新——BitLinear层的示例！

import torch

from bitnet import BitLinear

# 输入
x = torch.randn(10, 1000, 512)

# BitLinear层
layer = BitLinear(512, 400)

# 输出
y = layer(x)

print(y)

BitLinearNew

import torch
from bitnet import BitLinearNew

# 创建一个形状为(16, 10)的随机张量
x = torch.randn(16, 1000, 512)

# 创建一个输入大小为512、输出大小为20、具有2个组的BitLinearNew实例
layer = BitLinearNew(
    512,
    20,
)

# 将输入x通过BitLinearNew层进行前向传播
output = layer(x)

# 打印输出张量及其形状
print(output)
print(output.shape)

---

BitNetTransformer

• 完整实现了图示中的Transformer结构，包含多头注意力和BitFeedforward层
• 不仅可用于文本处理，还可用于图像，甚至视频或音频处理
• 全面支持残差连接和跳跃连接，以促进梯度流动

# 导入必要的库
import torch
from bitnet import BitNetTransformer

# 创建一个随机整数张量
x = torch.randint(0, 20000, (1, 1024))

# 初始化BitNetTransformer模型
bitnet = BitNetTransformer(
    num_tokens=20000,  # 输入中唯一标记的数量
    dim=1024,  # 输入和输出嵌入的维度
    depth=6,  # Transformer层数
    heads=8,  # 注意力头数
    ff_mult=4,  # 前馈网络中隐藏层维度的倍增因子
)

# 将张量通过Transformer模型传递
logits = bitnet(x)

# 打印输出张量的形状
print(logits)

BitAttention

该注意力机制已被修改为使用BitLinear代替默认的线性投影。同时，它采用了多组查询注意力而非常规多头注意力，以实现更快的解码速度和更长的上下文处理能力。

import torch
from bitnet import BitMGQA

# 创建一个形状为(1, 10, 512)的随机张量
x = torch.randn(1, 10, 512)

# 创建一个输入大小为512、具有8个注意力头和4层的BitMGQA模型
gqa = BitMGQA(512, 8, 4)

# 将输入张量通过BitMGQA模型传递，并获取输出及注意力权重
out, _ = gqa(x, x, x, need_weights=True)

# 打印输出张量和注意力张量的形状
print(out)

BitFeedForward

• 如图所示的前馈网络，包含BitLinear和GELU：
• 线性层 -> GELU -> 线性层
• 你还可以添加Dropout、层归一化或其他层，以构建更强大的前馈网络。

import torch
from bitnet import BitFeedForward

# 创建一个形状为(10, 512)的随机输入张量
x = torch.randn(10, 512)

# 创建一个BitFeedForward实例，参数如下：
# - input_dim: 512
# - hidden_dim: 512
# - num_layers: 4
# - swish: True（使用Swish激活函数）
# - post_act_ln: True（在每次激活后应用层归一化）
# - dropout: 0.1（以0.1的概率应用Dropout）
ff = BitFeedForward(512, 512, 4, swish=True, post_act_ln=True, dropout=0.1)

# 将输入张量x通过BitFeedForward网络处理
y = ff(x)

# 打印输出张量的形状
print(y)  # torch.Size([10, 512])

推理

from bitnet import BitNetInference

bitnet = BitNetInference()
bitnet.load_model("../model_checkpoint.pth")  # 下载模型
output_str = bitnet.generate("那只狗跳过了 ", 512)
print(output_str)

Hugging Face使用

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from bitnet import replace_linears_in_hf

# 从Hugging Face的Transformers库加载模型
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 将线性层替换为BitLinear
replace_linears_in_hf(model)

# 示例待分类文本
text = "将此处替换为你的文本"
inputs = tokenizer(
    text, return_tensors="pt", padding=True, truncation=True, max_length=512
)

# 进行推理
model.eval()  # 将模型设置为评估模式
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    print(predictions)

# 处理预测结果
predicted_class_id = predictions.argmax().item()
print(f"预测的类别ID：{predicted_class_id}")

# 如果你知道分类标签，可以将预测的类别ID映射到具体的标签上
# labels = ["标签1", "标签2", ...]  # 定义与模型类别对应的标签
# print(f"预测的标签：{labels[predicted_class_id]}")

替换PyTorch模型中的线性层

import torch
from torch import nn
from bitnet import replace_linears_in_pytorch_model

# 定义一个简单模型
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
)

print("替换前:")
print(model)

# 将 nn.Linear 替换为 BitLinear
replace_linears_in_pytorch_model(model)

print("替换后:")
print(model)

# 现在可以使用该模型进行训练或推理
# 例如，将一个随机输入传递给模型
input = torch.randn(1, 10)
output = model(input)

优化的 CUDA 内核

python setup.py build_ext --inplace

import torch
import gemm_lowbit_ext  # 导入编译后的模块

# 示例用法
a = torch.randn(10, 20, dtype=torch.half, device='cuda')  # 示例张量
b = torch.randn(20, 30, dtype=torch.half, device='cuda')  # 示例张量
c = torch.empty(10, 30, dtype=torch.half, device='cuda')  # 输出张量

w_scale = 1.0  # 示例缩放因子
x_scale = 1.0  # 示例缩放因子

# 调用自定义的 CUDA GEMM 操作
gemm_lowbit_ext.gemm_lowbit(a, b, c, w_scale, x_scale)

print(c)  # 查看结果

BitLora

BitLora 的实现！

import torch
from bitnet import BitLora

# 随机文本张量
x = torch.randn(1, 12, 200)

# 创建 BitLora 模型实例
model = BitLora(in_features=200, out_features=200, rank=4, lora_alpha=1)

# 进行前向传播
out = model(x)

# 打印输出张量的形状
print(out.shape)

BitMamba

import torch
from bitnet import BitMamba

# 创建一个大小为 (2, 10)、值在 0 到 100 之间的随机张量
x = torch.randint(0, 100, (2, 10))

# 创建一个 BitMamba 模型实例，输入大小为 512，隐藏层大小为 100，输出大小为 10，深度为 6
model = BitMamba(512, 100, 10, 6, return_tokens=True)

# 将输入张量传递给模型并获取输出
output = model(x)

# 打印输出张量
print(output)

# 打印输出张量的形状
print(output.shape)

BitMoE

import torch
from bitnet.bit_moe import BitMoE

# 创建输入张量
x = torch.randn(2, 4, 8)

# 创建一个指定输入和输出维度的 BitMoE 模型
model = BitMoE(8, 4, 2)

# 对模型进行前向传播
output = model(x)

# 打印输出
print(output)

1 位视觉 Transformer

这个想法突然冒出来，但看起来非常有趣，因为你可以利用 bitlinear 来处理视觉任务，实现超高压缩。如果能编写一个脚本，我们愿意提供计算资源来在 ImageNet 上训练它。接下来的阶段就是训练一个视觉-语言联合模型，比如 gpt-4o。

import torch
from bitnet import OneBitViT

# 创建 OneBitViT 模型实例
v = OneBitViT(
    image_size=256,
    patch_size=32,
    num_classes=1000,
    dim=1024,
    depth=6,
    heads=16,
    mlp_dim=2048,
)

# 生成一个随机图像张量
img = torch.randn(1, 3, 256, 256)

# 将图像传递给 OneBitViT 模型以获得预测
preds = v(img)  # (1, 1000)

# 打印预测结果
print(preds)

许可证

MIT

引用

@misc{2310.11453,
Author = {Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Huaijie Wang and Lingxiao Ma and Fan Yang and Ruiping Wang and Yi Wu and Furu Wei},
Title = {BitNet: Scaling 1-bit Transformers for Large Language Models},
Year = {2023},
Eprint = {arXiv:2310.11453},
}

待办事项

• 再次检查 BitLinear 的实现，确保其完全符合论文中的描述
• 实现 BitNetTransformer 的训练脚本
• 在 Enwiki8 数据集上进行训练，复制 Lucidrains 仓库中的代码和数据
• 基准测试性能
• 研究针对非可微反向传播的直通估计器
• 实现 BitFeedForward
• 清理代码库
• 为每个模块添加单元测试
• 实现来自 论文 的新版本 BitNet1.5b
• 在 CUDA 中实现 BitNet15b
• 实现低比特 GEMM CUDA 内核

BitNet 快速上手指南

BitNet 是一个基于 PyTorch 的开源项目，实现了论文《BitNet: Scaling 1-bit Transformers for Large Language Models》中的 1-bit 线性层（BitLinear）及相关模型架构。它通过将 Transformer 中的标准线性投影替换为 1-bit 量化版本，旨在大幅降低大语言模型的存储和计算成本。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux 或 macOS (Windows 需配置相应的 CUDA 环境)
• Python: 3.8 或更高版本
• 深度学习框架: PyTorch (建议安装支持 CUDA 的版本以启用 GPU 加速)
• 硬件: 推荐使用 NVIDIA GPU 以获得最佳性能（部分算子依赖 CUDA）

前置依赖安装： 如果尚未安装 PyTorch，请先访问 PyTorch 官网 获取适合您环境的安装命令。国内用户可使用清华源加速安装：

pip3 install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

您可以直接通过 pip 安装 BitNet 库：

pip3 install bitnet

可选：编译优化的 CUDA 内核 如果您需要使用自定义的低精度 GEMM CUDA 算子以获得极致性能，请在克隆源码后执行以下命令进行本地编译：

python setup.py build_ext --inplace

基本使用

BitNet 的核心创新在于 BitLinear 层，它可以无缝替换 PyTorch 中的 nn.Linear。以下是几种最常用的场景。

1. 使用 BitLinear 层

这是最基础的用法，直接实例化 BitLinear 模块并进行前向传播。

import torch
from bitnet import BitLinear

# 输入张量 (Batch, Sequence, Dimension)
x = torch.randn(10, 1000, 512)

# 创建 BitLinear 层：输入维度 512 -> 输出维度 400
layer = BitLinear(512, 400)

# 前向传播
y = layer(x)

print(y.shape)  # 输出形状：torch.Size([10, 1000, 400])

2. 快速替换现有 PyTorch 模型

如果您有一个已经定义好的 PyTorch 模型，可以使用工具函数一键将其中所有的 nn.Linear 替换为 BitLinear。

注意：替换后的模型需要从头训练（Train from scratch）或进行微调，直接加载原模型权重无法正常工作。

import torch
from torch import nn
from bitnet import replace_linears_in_pytorch_model

# 定义一个普通的 PyTorch 模型
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
)

print("替换前:")
print(model)

# 执行替换
replace_linears_in_pytorch_model(model)

print("替换后:")
print(model)

# 测试推理
input_tensor = torch.randn(1, 10)
output = model(input_tensor)
print(output.shape)

3. 使用完整的 BitNet Transformer

BitNet 提供了完整的 Transformer 实现，包含多头注意力机制（已优化为 Multi-Grouped Query Attention）和位量化前馈网络。

import torch
from bitnet import BitNetTransformer

# 模拟输入 token IDs (Batch=1, Seq_Len=1024)
x = torch.randint(0, 20000, (1, 1024))

# 初始化模型
bitnet = BitNetTransformer(
    num_tokens=20000,  # 词表大小
    dim=1024,          # 嵌入维度
    depth=6,           # Transformer 层数
    heads=8,           # 注意力头数
    ff_mult=4,         # 前馈网络扩展倍数
)

# 前向传播
logits = bitnet(x)

print(logits.shape)

4. Hugging Face 模型集成

您可以轻松地将 Hugging Face transformers 库中的模型转换为 BitNet 架构。

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from bitnet import replace_linears_in_hf

# 加载预训练模型 (例如 BERT)
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 将模型中的 Linear 层替换为 BitLinear
replace_linears_in_hf(model)

# 准备输入数据
text = "Replace this with your text"
inputs = tokenizer(
    text, return_tensors="pt", padding=True, truncation=True, max_length=512
)

# 推理
model.eval()
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    print(predictions)

---

重要提示：由于 BitNet 采用了 1-bit 量化策略，其权重分布与传统浮点模型完全不同。切勿直接将传统模型的权重加载到替换后的 BitNet 模型中，必须重新训练或从头微调。