BitBLAS技术文档

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)

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

项目API使用文档

BitBLAS的核心在于Matmul和Linear两个类，它们提供了灵活的数据类型支持和高效计算功能：

• Matmul: 用户可以通过指定权重($wdtype$)、激活($adtype$)和输出($cdtype$)的数据类型，执行自定义的混合精度矩阵乘法。
• Linear: 提供类似PyTorch线性层的操作接口，自动管理数据类型的转换，便于集成到现有的DNN框架中。

项目安装方式

除上述通过pip的简单安装外，对于特殊需求或不兼容环境，推荐查看项目GitHub页面下的源码构建指南，遵循文档步骤以适应特定的系统配置。

BitBLAS以其高性能的混合精度运算能力，旨在为深度学习推理和训练带来显著的效率提升，尤其是对资源敏感和追求速度的应用场景而言，是不可或缺的工具之一。

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