BitBLAS技术文档

张旦宪Sacha

于 2024-10-18 12:06:07 发布

阅读量639

点赞数 28

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_01241/article/details/143045778

BitBLAS技术文档

BitBLAS BitBLAS is a library to support mixed-precision matrix multiplications, especially for quantized LLM deployment. 项目地址: https://gitcode.com/gh_mirrors/bi/BitBLAS

BitBLAS 是一个专为GPU设计的混合精度BLAS运算库，优化了如 $W_{wdtype}A_{adtype}$ 的矩阵乘法操作，典型应用在深度神经网络（DNN）模型的部署中，尤其适用于大型语言模型（LLMs）的高效量化处理，比如在GPTQ中的 $W_{UINT4}A_{FP16}$、BitDistiller的 $W_{INT2}A_{FP16}$ 以及BitNet-b1.58的 $W_{INT2}A_{INT8}$ 等场景。该库基于“Ladder”论文的技术，该论文在OSDI'24会议上发表。

安装指南

使用pip安装

确保您的系统满足以下前提条件：

操作系统: Ubuntu 20.04 或更高版本
Python版本: 3.8及以上
CUDA版本: 11.0以上

通过pip直接安装BitBLAS的最新版，执行以下命令：

pip install bitblas

安装后，可验证是否成功：

python -c "import bitblas; print(bitblas.__version__)"

注意事项: 目前仅提供Ubuntu 20.04平台上的whl文件，并且支持CUDA≥11.0及Python≥3.8。其他环境需从源码编译安装，请参考安装文档。

项目使用说明

BitBLAS提供了两种Python API来实现混合精度矩阵乘法：

bitblas.Matmul: 实现了$W_{wdtype}A_{adtype}→C_{cdtype}$的矩阵乘法，其中$wdtype$、$adtype$表示权重和激活数据类型，$cdtype$表示结果的数据类型。
bitblas.Linear: 类似于PyTorch的nn.Linear模块，支持混合精度线性操作。

示例：$W_{INT4}A_{FP16}$混合精度矩阵乘法

以执行$W_{INT4}A_{FP16}$混合精度矩阵乘法为例，代码展示输入矩阵的创建、权重的量化处理，以及利用bitblas.Matmul进行计算的过程，并对比传统方法得到的结果验证准确性。

# 假设A为FP16矩阵，W为经过量化处理至INT4的矩阵
import torch
import bitblas

A = torch.randn(100, 50).half()  # FP16激活矩阵
W = torch.randn(30, 50).half().quantize_per_tensor(scale=1.0/127.0, zero_point=0, dtype=torch.int4)  # 权重矩阵量化至INT4

# 初始化并执行混合精度矩阵乘法
output = bitblas.Matmul('INT4', 'FP16', 'FP16')(A, W)
print("Output Shape:", output.shape)

请注意，这里示例中的量化过程简化处理，实际使用时需结合具体量化策略。