[论文品鉴] 1-bit LLM

今天看了一篇4月最新论文 <BitNet b1.58 2B4T Technical Report>，觉得有点意思，决定写点东西沉淀一下。

1-bit LLM

首先要整明白的就是1-bit LLM，这年头都在搞量化，但从没听说过能量化到1-bit，翻了一下arxiv才发现也不是啥新东西，24年的论文 <The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits> 就有过相关的论证，下面简单介绍下。

1-bit LLM的实现，主要靠的是极致的量化和新的架构设计，把传统浮点参数量化为极致的bit表示，如三值 { − 1 ， 0 ， 1 } \{-1，0，1\} {−1，0，1}。

说到这肯定有会疑问，1 bit怎么可能表示3个值？

正常情况下1 bit就是 { 0 , 1 } \{0,1\} {0,1}，这是底层操作系统二进制决定的，所以如果想要表示3个值，那么理论上至少需要 $lo{g}_{2}3\approx 1.58bits$ 才行。因为本质上三值法并没有创建新的底层系统，还是二进制，也就只有 { 0 , 1 } \{0,1\} {0,1}，所以k 个 bit 可以表示出 $2^k$ 种组合，那么对于3种数值就是$2^k=3$ 得出 $k=lo{g}_{2}3\approx 1.58$ 它是个平均值。

掩码存储

另一方面，三值在底层存储的时候，采用了两种方式实现：

• 将多个参数打包到整数中，例如8个bit存5个值
• 特殊编码和掩码

这里详细讲解一下有趣的特殊编码和掩码设计：

• 特殊编码：1 bit 用于存储符号，所以通过这1 bit可以表示出三值中的两个，也就是 { − 1 , + 1 } \{-1,+1\} {−1,+1}
• 掩码：还剩的那个 $0$ 咋整，这也正是1 bit LLM实现的关键，$0$ 靠的正是掩码实现的，而非存储实际数值。

整体上分三步走：

1. 三值量化：例如对于4个参数 $[+1,2,0,-0.8,0]$ ，首先三值离散量化为 $[+1,0,-1,0]$
2. 生成掩码：这一步只需要标记出非零的位置，算是一种位置编码，$Bitmask=[1,0,1,0]$，底层存储二进制 $1010$
3. 符号存储：这一步只需要存储非零元素的符号，$Symbols=[1,0]$，其中$+1$的符号表示为$1$，$-1$的符号表示为$0$，底层存储二进制 $10$

所以通过以上新的存储设计，实现了4个参数6bit存储，也就是对于$N$个参数，需要$N$个位置编码与$0.5N$个符号编码（这里假设50%是0）那么就有$1.5N$，所以每个参数平均需要$\frac{1.5N}{N}=1.5bits$

量化函数

$W_{quant}=Sign(W)\cdot \gamma +ϵ$ ，动态缩放因子和残差项不必说，只需要关注$sign()$

因为三值离散量化，也就是把浮点数W量化为$[+1,0,-1,0]$ ，那么对于一个阈值Δ，则有：$$W_{quant}=\{\begin{array}{ll}+1& \text{if w > Δ}\\ 0& \text{if -Δ ≤ w ≤ Δ}\\ -1& \text{if w < -Δ}\end{array}$$
正好$Sign()$完美符合条件，对于小于0的数结果为-1，大于0的数结果为1，0的结果为0。

BitNet b1.58 2B-4T

上面说的都是理论，微软基于此开源了，通过4 trillon个token训练出来的2B模型 BitNet b1.58 2B4T的LLM，Hugging face上已经可以看到了，它主要有2点核心设计：

1. BitLinear线性层 替代了全精度线性层：也就是三值量化
2. 硬件优化：有专用的推理框架bitnet.cpp，通过查表法和smid加速了推理

使用方式的选择也比较全：python、cpp、docker等，同时支持cpu、gpu

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' # vpn慢可以换成国内镜像

!pip install git+https://github.com/shumingma/transformers.git # 亲测，目前还必须使用这个版本的库

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "microsoft/bitnet-b1.58-2B-4T"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16
)

messages = [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "How are you?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
chat_input = tokenizer(prompt, return_tensors="pt").to(model.device)

chat_outputs = model.generate(**chat_input, max_new_tokens=50)
response = tokenizer.decode(chat_outputs[0][chat_input['input_ids'].shape[-1]:], skip_special_tokens=True) # Decode only the response part
print("\nAssistant Response:", response)