필사 모드: BitNet 论文解析：1-Bit LLM 的时代 — 从三值权重到 CPU 推理

W_b = Sign(W) = +1  (if W >= 0)
                -1  (if W < 0)

alpha = (1/nm) * sum(|W_ij|)   # 缩放因子

import torch
import torch.nn as nn
import torch.nn.functional as F

class BitLinear_v1(nn.Module):
    """BitNet v1 的 BitLinear 实现（教学用简化版本）"""
    def __init__(self, in_features, out_features, activation_bits=8):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.activation_bits = activation_bits
        self.Qb = 2 ** (activation_bits - 1)
        # 实数权重（训练时更新）
        self.weight = nn.Parameter(torch.randn(out_features, in_features))

    def ste_binarize(self, w):
        """使用 Straight-Through Estimator 进行二值化"""
        # 前向传播：应用 sign 函数
        # 反向传播：梯度直接通过（STE）
        w_bin = w.sign()
        # STE：用 detach() 阻断 sign 的梯度，保留原始 w 的梯度
        return w + (w_bin - w).detach()

    def activation_quant(self, x):
        """激活值 absmax 量化"""
        gamma = x.abs().max()
        x_quant = torch.clamp(x * self.Qb / (gamma + 1e-5), -self.Qb, self.Qb - 1)
        return x_quant, gamma

    def forward(self, x):
        # 权重二值化 + 缩放因子
        w_bin = self.ste_binarize(self.weight)
        alpha = self.weight.abs().mean()

        # 激活值量化
        x_quant, gamma = self.activation_quant(x)

        # 整数矩阵运算（用加减法代替乘法）
        output = F.linear(x_quant, w_bin)

        # 反量化：恢复尺度
        output = output * (alpha * gamma) / self.Qb
        return output

import torch
import torch.nn as nn
import torch.nn.functional as F

def weight_quant_ternary(w):
    """BitNet b1.58 三值权重量化（基于 absmean）"""
    # 缩放因子：权重绝对值的均值
    gamma = w.abs().mean()
    # 缩放后通过 round、clamp 限制为 {-1, 0, +1}
    w_scaled = w / (gamma + 1e-5)
    w_ternary = torch.clamp(torch.round(w_scaled), -1, 1)
    # STE：前向传播使用量化值，反向传播使用原始梯度
    return w + (w_ternary - w).detach(), gamma

class BitLinear_b158(nn.Module):
    """BitNet b1.58 的 BitLinear 实现"""
    def __init__(self, in_features, out_features, activation_bits=8):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.02)
        self.activation_bits = activation_bits
        self.Qb = 2 ** (activation_bits - 1)
        # 在激活前应用 RMSNorm
        self.norm = nn.RMSNorm(in_features)

    def activation_quant(self, x):
        """激活值 absmax 量化（按 token）"""
        gamma = x.abs().max(dim=-1, keepdim=True).values
        x_q = torch.clamp(
            torch.round(x * self.Qb / (gamma + 1e-5)),
            -self.Qb, self.Qb - 1
        )
        return x_q, gamma

    def forward(self, x):
        # 激活值归一化
        x = self.norm(x)

        # 三值权重量化
        w_q, w_scale = weight_quant_ternary(self.weight)

        # 激活值量化
        x_q, x_scale = self.activation_quant(x)

        # 整数运算：乘法被加法/减法/跳过所替代
        output = F.linear(x_q, w_q)

        # 反量化
        output = output * (w_scale * x_scale) / self.Qb
        return output

# BitNet b1.58 2B4T 模型下载与 bitnet.cpp 推理环境配置
# 1. 克隆仓库
git clone --recursive https://github.com/microsoft/BitNet.git
cd BitNet

# 2. 安装依赖（推荐使用 conda 环境）
conda create -n bitnet python=3.11 -y
conda activate bitnet
pip install -r requirements.txt

# 3. 下载模型并构建推理引擎（一次性完成）
python setup_env.py --hf-repo microsoft/BitNet-b1.58-2B-4T-gguf \
    -q i2_s \
    --quant-embd

# 4. 执行推理
python run_inference.py -m models/BitNet-b1.58-2B-4T-gguf/ggml-model-i2_s.gguf \
    -p "Microsoft Research recently released" \
    -n 128 \
    -t 4 \
    --temp 0.7

import numpy as np

def tl2_lookup_simulation(activations, ternary_weights_packed):
    """基于 TL2 查找表的三值矩阵-向量乘模拟"""
    # 针对两个连续权重 (w0, w1) 组合的全部 9 种可能
    # 预先计算与激活值 (a0, a1) 的内积结果
    # 用查表替代 w0*a0 + w1*a1
    #
    # 编码：w=(-1,0,1) -> (0,1,2)，组合 idx = w0_enc * 3 + w1_enc
    # idx=0: (-1,-1) -> -(a0+a1)
    # idx=1: (-1, 0) -> -a0
    # idx=2: (-1,+1) -> -a0+a1
    # idx=3: ( 0,-1) -> -a1
    # idx=4: ( 0, 0) -> 0
    # idx=5: ( 0,+1) -> a1
    # idx=6: (+1,-1) -> a0-a1
    # idx=7: (+1, 0) -> a0
    # idx=8: (+1,+1) -> a0+a1

    n = len(activations)
    result = 0
    for i in range(0, n, 2):
        a0, a1 = activations[i], activations[i+1]
        # 生成查找表（实际实现中会加载到 SIMD 寄存器）
        lut = [
            -(a0 + a1), -a0, -a0 + a1,
            -a1, 0, a1,
            a0 - a1, a0, a0 + a1
        ]
        # 从打包索引中提取组合
        idx = ternary_weights_packed[i // 2]  # 0~8
        result += lut[idx]
    return result

# 验证
np.random.seed(42)
acts = np.random.randn(8).astype(np.float32)
# 三值权重：[-1, 1, 0, 1, -1, 0, 1, -1]
weights = np.array([-1, 1, 0, 1, -1, 0, 1, -1])
# 生成打包索引
packed = []
for i in range(0, 8, 2):
    w0_enc = weights[i] + 1  # {-1,0,1} -> {0,1,2}
    w1_enc = weights[i+1] + 1
    packed.append(w0_enc * 3 + w1_enc)

ref = np.dot(acts, weights)
tl2 = tl2_lookup_simulation(acts, packed)
print(f"Reference dot product: {ref:.6f}")
print(f"TL2 lookup result:     {tl2:.6f}")
print(f"Match: {np.isclose(ref, tl2)}")

# 执行 bitnet.cpp 性能基准测试
cd BitNet

# 单线程基准测试
python utils/benchmark.py \
    -m models/BitNet-b1.58-2B-4T-gguf/ggml-model-i2_s.gguf \
    -n 512 \
    -p 256 \
    --threads 1

# 多线程基准测试（匹配物理核心数）
python utils/benchmark.py \
    -m models/BitNet-b1.58-2B-4T-gguf/ggml-model-i2_s.gguf \
    -n 512 \
    -p 256 \
    --threads 4

# 测量不同提示长度下的性能
for prompt_len in 64 128 256 512 1024; do
    echo "=== Prompt length: $prompt_len ==="
    python utils/benchmark.py \
        -m models/BitNet-b1.58-2B-4T-gguf/ggml-model-i2_s.gguf \
        -n 128 \
        -p $prompt_len \
        --threads 4
done

# 使用 BitNet 模型的简单文本生成流水线示例
# （实际部署时请使用 bitnet.cpp 的 C++ API）
import subprocess
import json
import sys

class BitNetInference:
    """基于 bitnet.cpp 的推理封装类"""
    def __init__(self, model_path, n_threads=4):
        self.model_path = model_path
        self.n_threads = n_threads
        self.binary = "./build/bin/llama-cli"  # bitnet.cpp 构建生成的二进制文件

    def generate(self, prompt, max_tokens=128, temperature=0.7, top_p=0.9):
        """文本生成"""
        cmd = [
            self.binary,
            "-m", self.model_path,
            "-p", prompt,
            "-n", str(max_tokens),
            "-t", str(self.n_threads),
            "--temp", str(temperature),
            "--top-p", str(top_p),
            "--no-display-prompt"
        ]
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=120
            )
            if result.returncode != 0:
                raise RuntimeError(f"Inference failed: {result.stderr}")
            return result.stdout.strip()
        except subprocess.TimeoutExpired:
            raise TimeoutError("Inference timed out after 120 seconds")

    def benchmark(self, prompt_lengths=[64, 128, 256, 512]):
        """测量不同提示长度下的性能"""
        results = {}
        for length in prompt_lengths:
            prompt = "A " * length  # 占位提示词
            cmd = [
                self.binary,
                "-m", self.model_path,
                "-p", prompt,
                "-n", "1",  # 仅生成 1 个 token 以测量 prefill 速度
                "-t", str(self.n_threads),
                "--no-display-prompt"
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            # 从 stderr 解析性能指标
            for line in result.stderr.split('\n'):
                if 'eval time' in line:
                    # 提取 token/秒
                    parts = line.split()
                    for i, p in enumerate(parts):
                        if p == 'token/s)':
                            tok_per_sec = float(parts[i-1].strip('('))
                            results[length] = tok_per_sec
        return results

# 使用示例
if __name__ == "__main__":
    model_path = "models/BitNet-b1.58-2B-4T-gguf/ggml-model-i2_s.gguf"
    engine = BitNetInference(model_path, n_threads=4)

    # 文本生成
    output = engine.generate(
        "The key advantages of 1-bit LLMs are:",
        max_tokens=200,
        temperature=0.8
    )
    print(output)

# 问题场景：CMake 版本低于 3.22 时构建失败
# 错误信息：
# CMake Error at CMakeLists.txt:1:
#   CMake 3.22 or higher is required. You are running version 3.16.3

# 解决方法 1：升级 CMake（Ubuntu）
sudo apt remove cmake
pip install cmake --upgrade
# 或者
sudo snap install cmake --classic

# 解决方法 2：在 conda 环境中安装 CMake
conda install -c conda-forge cmake>=3.22

# 解决方法 3：从源码构建
wget https://github.com/Kitware/CMake/releases/download/v3.28.3/cmake-3.28.3.tar.gz
tar xzf cmake-3.28.3.tar.gz
cd cmake-3.28.3
./bootstrap && make -j$(nproc) && sudo make install

# 重试构建
cd BitNet
python setup_env.py --hf-repo microsoft/BitNet-b1.58-2B-4T-gguf \
    -q i2_s --quant-embd