T-MAC

BitNet 在 T-MAC（基于查找表）与 llama.cpp（基于反量化）上的对比

新闻

• 2024年8月14日 🚀：T-MAC GEMM (N>1) 内核现已集成到 llama.cpp 中，以加速预填充。

• 2024年8月6日 🚀：支持 GPTQ 格式的 1/2/3/4 位量化 Llama 模型。使用 EfficientQAT 发布的预训练模型进行测试。

• 2024年7月27日 ✨：我们注意到在最新的骁龙 X Elite 芯片组上，T-MAC 在token生成速度上甚至比 NPU 更快！查看与 NPU 比较获取更多详情。

• 2024年7月23日 🚀🚀：我们已经支持通过 T-MAC 执行任何 GPTQ 格式的 2 位量化 Llama 模型！使用 EfficientQAT 发布的预训练模型进行测试。

• 2024年7月22日 🚀🚀：我们已经为 ARM 架构的 Windows 添加了原生部署支持。T-MAC 在 Surface Laptop 7 上展示了显著的 5 倍速度提升。

简介

T-MAC 是一个内核库，通过使用查找表直接支持混合精度矩阵乘法（int1/2/3/4 x int8/fp16/fp32），无需反量化。T-MAC 旨在提升 CPU 上低比特 LLM 推理的性能。T-MAC 已经支持各种低比特模型，包括来自 GPTQ/gguf 的 W4A16、来自 BitDistiller/EfficientQAT 的 W2A16 以及来自 BitNet 的 W1(.58)A8，适用于配备 ARM/Intel CPU 的 OSX/Linux/Windows 系统。

对于 3B BitNet，T-MAC 在 Surface Laptop 7 上单核可达到 20 tokens/秒的 token 生成吞吐量，四核可达到 48 tokens/秒，相比于最先进的 CPU 低比特框架（llama.cpp）实现了 4~5 倍的加速。T-MAC 甚至可以在树莓派 5 等性能较低的设备上达到 11 tokens/秒。

端到端加速

我们在五种不同的设备上评估了不同模型的 token 生成性能：Surface Laptop 7、Apple M2-Ultra、Jetson AGX Orin、树莓派 5 和 Surface Book 3。查看数据表获取更多详情。

我们使用 T-MAC 2 位和 llama.cpp Q2_K 评估 BitNet-3B 和 Llama-2-7B (W2)，使用 T-MAC 4 位和 llama.cpp Q4_0 评估 Llama-2-7B (W4)。

除了提供显著的加速外，T-MAC 还可以使用更少的 CPU 核心达到相同的性能。例如，要达到 40 tokens/秒（远超人类阅读速度的吞吐量），T-MAC 仅需要 2 个核心，而 llama.cpp 需要 8 个核心。在 Jetson AGX Orin 上，要达到 10 tokens/秒（已经满足人类阅读速度的吞吐量），T-MAC 仅需要 2 个核心，而 llama.cpp 使用了全部 12 个核心。T-MAC 可以在树莓派 5 等配备较少 CPU 核心的低功率设备上满足实时需求。通过使用更少的核心，T-MAC 可以为其他应用保留计算资源，并显著降低功耗和能耗，这对边缘设备至关重要。

T-MAC 在单线程上实现显著加速，并消耗更少的 CPU 核心达到相同的吞吐量

T-MAC 的吞吐量是在不使用快速聚合的情况下获得的。用户可以通过 -fa 开启快速聚合，以获得额外 10%~20% 的加速。

预填充加速

待办：添加更多结果

我们对比了 Surface Laptop 7 上 Llama-2-7b (W2) 的预填充吞吐量（input_len=256），基准为：

• llama.cpp：llama.cpp 优化的基于反量化的低比特内核
• llama.cpp (OpenBLAS)：llama.cpp OpenBLAS 后端

内核级加速

我们的 GEMM 内核在 CPU 上展示了优于最先进的低比特 GEMM 的性能。下图显示了 token 生成过程中 llama-7b 内核相比 llama.cpp 的加速比（单线程）：

llama.cpp 没有提供 1 位内核实现，但我们可以从 2 位推断，因为根据 2/3/4 位的结果，它不会带来额外的加速。

由于计算成本的降低，T-MAC 可以在多批次（N>1）GEMM 中实现显著加速，这确保了在提示评估和多批次 token 生成方面的出色性能。下图显示了与使用 OpenBLAS 后端的 llama.cpp 相比的加速比（单线程）：

M2-Ultra 是个例外，因为它配备了专门设计的 AMX 协处理器 来加速多批次 GEMM。然而，T-MAC 在 2 位时仍然可以达到comparable的性能。

节能和降低功耗

通过用查表指令替代繁重的融合乘加指令，T-MAC 显著降低了功耗。结合加速效果，T-MAC 最终导致总能耗的大幅降低。

M2-Ultra 上三种模型的多线程功耗/能耗，M1: Llama-2-7B (W4)，M2: Llama-2-7B (W2) 和 M3: BitNet-3B

数据使用 powermetrics 采样。

与 NPU 比较

在最新的骁龙 X Elite 芯片组上，通过 T-MAC 的 CPU 相比通过高通骁龙神经处理引擎（NPE）的 NPU 实现了更好的性能。

在部署 llama-2-7b-4bit 模型时，NPU 只能生成 10.4 tokens/秒（根据此处发布的数据），而使用 T-MAC 的 CPU 可以用两个核心达到 12.6 tokens/秒，甚至高达 22 tokens/秒。考虑到 T-MAC 的计算性能可以随着比特数的减少而线性提升（这在基于反量化的 GPU 和 NPU 上无法观察到），T-MAC 甚至可以在 2 位时用单核 CPU 匹配 NPU 的性能。

为了公平比较，我们将设置与NPU对齐，包括1024的输入长度和1024的输出长度。尽管Qualcomm部署了3.6GB的模型，由于我们的token嵌入保持未量化，我们部署的模型略大，为3.7GB。

与CUDA GPU比较

T-MAC在Jetson AGX Orin上实现了与CUDA GPU相当的2比特mpGEMM性能。虽然CUDA GPU在执行mpGEMM以外的内核时表现更佳，使得T-MAC (CPU)的端到端性能略慢，但T-MAC可以在功耗和能耗方面带来显著节省。

NVIDIA Jetson AGX Orin上Llama-2-7B (W2)的吞吐量/功耗/能耗比较（CPU线程数为12）

数据使用jetson-stats在MAXN功耗模式下采样。

安装

要求

• Python (TVM需要3.8版本)
• virtualenv
• cmake>=3.22

brew install cmake zlib libomp

git clone --recursive https://github.com/microsoft/T-MAC.git
# 在virtualenv中
pip install . -v  # 或 pip install -e . -v
source build/t-mac-envs.sh

sudo apt install build-essential libtinfo-dev zlib1g-dev libzstd-dev libxml2-dev

git clone --recursive https://github.com/microsoft/T-MAC.git
# 在virtualenv中
pip install . -v  # 或 pip install -e . -v
source build/t-mac-envs.sh

git clone --recursive https://github.com/microsoft/T-MAC.git
cd T-MAC
conda env create --file conda\tvm-build-environment.yaml
conda activate tvm-build

cd 3rdparty\tvm
mkdir build
cp cmake\config.cmake build

cd build
cmake .. -A x64
cmake --build . --config Release -- /m

cd ..\..\..\  # 回到项目根目录
$env:MANUAL_BUILD = "1"
$env:PYTHONPATH = "$pwd\3rdparty\tvm\python"
pip install . -v  # 或 pip install -e . -v

git clone --recursive https://github.com/microsoft/T-MAC.git
cd T-MAC
conda env create --file conda\tvm-build-environment.yaml
conda activate tvm-build

cd 3rdparty\tvm
mkdir build
cp cmake\config.cmake build

cd build
cmake .. -A x64  # 以x64构建TVM，因为Python和依赖项是x64的
cmake --build . --config Release -- /m

conda activate tvm-build
conda uninstall cmake  # 防止与原生ARM64 cmake潜在冲突
cd ..\..\..\  # 回到项目根目录
$env:MANUAL_BUILD = "1"
$env:PYTHONPATH = "$pwd\3rdparty\tvm\python"
pip install wmi  # 用于在x86_64 python中检测原生ARM64 CPU
pip install . -v  # 或 pip install -e . -v

验证

之后，你可以通过以下方式验证安装：python -c "import t_mac; print(t_mac.__version__); from tvm.contrib.clang import find_clang; print(find_clang())"。

使用方法

目前，我们通过llama.cpp集成支持端到端推理。

我们提供了一个一体化脚本。使用以下命令调用：

pip install 3rdparty/llama.cpp/gguf-py
huggingface-cli download 1bitLLM/bitnet_b1_58-3B --local-dir ${model_dir}
python tools/run_pipeline.py -o ${model_dir}

我们还支持来自GPTQModel/EfficientQAT的GTPQ格式模型。尝试使用官方发布的EfficientQAT（GPTQ格式）Llama-3-8b-instruct-w2-g128:

huggingface-cli download ChenMnZ/Llama-3-8b-instruct-EfficientQAT-w2g128-GPTQ --local-dir ${model_dir}
python tools/run_pipeline.py -o ${model_dir} -m llama-3-8b-2bit

使用 -p 或 -s 参数选择要运行的步骤。使用 -u 参数来使用我们为ARM预构建的内核。

对于非对称量化，使用 --zero_point，这对大多数EfficientQAT模型是必需的（仅在Llama-3-8b-instruct-w4-g128/Llama-3-8b-instruct-w2-g128上验证过）。

输出示例：

运行步骤0：编译内核
  在/Users/user/jianyu/T-MAC/deploy目录下运行命令：
    python compile.py -o tuned -da -nt 4 -tb -gc -gs 128 -ags 64 -t -m hf-bitnet-3b -r
运行步骤1：构建T-MAC C++ CMakeFiles
  在/Users/user/jianyu/T-MAC/build目录下运行命令：
    cmake -DCMAKE_INSTALL_PREFIX=/Users/user/jianyu/T-MAC/install ..
运行步骤2：安装T-MAC C++
  在/Users/user/jianyu/T-MAC/build目录下运行命令：
    cmake --build . --target install --config Release
运行步骤3：将HF转换为GGUF
  在/Users/user/jianyu/T-MAC/3rdparty/llama.cpp目录下运行命令：
    python convert-hf-to-gguf-t-mac.py /Users/user/Downloads/test_models/hf-bitnet-3B --outtype i2 --outfile /Users/user/Downloads/test_models/hf-bitnet-3B/ggml-model.i2.gguf --kcfg /Users/user/jianyu/T-MAC/install/lib/kcfg.ini
运行步骤4：构建llama.cpp CMakeFiles
  在/Users/user/jianyu/T-MAC/3rdparty/llama.cpp/build目录下运行命令：
    cmake .. -DLLAMA_TMAC=ON -DCMAKE_PREFIX_PATH=/Users/user/jianyu/T-MAC/install/lib/cmake/t-mac -DCMAKE_BUILD_TYPE=Release -DLLAMA_LLAMAFILE_DEFAULT=OFF -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++
运行步骤5：构建llama.cpp
  在/Users/user/jianyu/T-MAC/3rdparty/llama.cpp/build目录下运行命令：
    cmake --build . --target main --config Release
运行步骤6：执行推理
  在/Users/user/jianyu/T-MAC/3rdparty/llama.cpp/build目录下运行命令：
    /Users/user/jianyu/T-MAC/3rdparty/llama.cpp/build/bin/main -m /Users/user/Downloads/test_models/hf-bitnet-3B/ggml-model.i2.gguf -n 128 -t 4 -p Microsoft Corporation is an American multinational corporation and technology company headquartered in Redmond, Washington. -b 1 -ngl 0 -c 2048
查看logs/2024-07-15-17-10-11.log获取推理输出

即将推出的功能

我们即将：

• 添加I4格式以简化4位模型的部署。
• 将T-MAC GEMM内核嵌入llama.cpp以加速预填充/提示。
• 通过LUTI4优化支持SME2的ARMv9 CPU

技术

LLM推理需要大量计算成本。低位量化是一种广泛采用的技术，但它引入了混合精度GEMM（mpGEMM）的挑战，硬件不直接支持这种运算，需要进行转换/反量化操作。

我们提出使用查找表（LUT）来支持mpGEMM。我们的方法包括以下关键技术：

1. 鉴于权重的低精度，我们将一位权重分组（例如，分成4组），预计算所有可能的部分和，然后使用LUT存储它们。
2. 我们采用移位和累加操作来支持从1到4的可扩展位数。
3. 在CPU上，我们利用tbl/pshuf指令进行快速表查找。
4. 我们将表大小从$2^n$减少到$2^{n-1}$，并结合符号位来加速LUT预计算。

我们的方法展现了几个显著特征：

1. T-MAC显示了FLOPs和推理延迟相对于位数的线性缩放比例。这与传统的基于转换的方法形成对比，后者在从4位降低到更低位时无法实现额外的加速。
2. T-MAC天然支持int1/2/3/4的按位计算，无需反量化。此外，它通过快速表查找和加法指令适应所有类型的激活（如fp8、fp16、int8），绕过了对支持不佳的融合乘加指令的需求。

引用

如果您觉得这个仓库有用，请使用以下BibTeX条目进行引用。

@misc{wei2024tmaccpurenaissancetable,
      title={T-MAC: CPU Renaissance via Table Lookup for Low-Bit LLM Deployment on Edge}, 
      author={Jianyu Wei and Shijie Cao and Ting Cao and Lingxiao Ma and Lei Wang and Yanyong Zhang and Mao Yang},
      year={2024},
      eprint={2407.00088},
      archivePrefix={arXiv},
      primaryClass={cs.DC},
      url={https://arxiv.org/abs/2407.00088}, 
}