LLMs

主题切换

自定义 CUDA函数的轻量级包装器 —— bitsandbytes 篇 ​

来源：AiGC面试宝典 发布时间：2024年01月27日

一、什么是 bitsandbytes？ ​

bitsandbytes 是自定义 CUDA 函数的轻量级包装器，特别是 8 比特优化器、矩阵乘法和量化函数。主要特征如下：

  • 具有混合精度分解的 8 比特矩阵乘法

    📝通俗解释：普通的矩阵乘法需要很高的计算精度（32比特），但这样很占内存。8比特矩阵乘法把计算精度降低到8比特（只有256个可能的值），同时通过混合精度分解技术保持计算结果的准确性。这样可以让大模型的推理速度更快，显存占用更少。

  • LLM.int8() 推理

    📝通俗解释：这是专门为大语言模型设计的8比特推理方法。传统的模型推理需要完整的16比特或32比特精度，非常消耗显存。LLM.int8() 可以把推理精度降到8比特，显存占用减少一半，同时几乎不影响答案的正确性。

  • 8 比特优化器：Adam、AdamW、RMSProp、LARS、LAMB、Lion（节省 75% 内存）

    📝通俗解释：训练大模型时，优化器需要保存很多中间状态（比如 momentum、variance 等），这些状态非常占内存。8比特优化器把这些状态压缩到8比特来存储，可以节省约75%的显存，这样就能用更少的GPU训练更大的模型。

  • 稳定的嵌入层：通过更好的初始化和标准化提高稳定性

    📝通俗解释：嵌入层是模型处理文字的第一步，如果初始化不好，模型训练容易出问题。这个功能提供了更好的初始化方法，让模型训练更稳定。

  • 8 比特量化：分位数、线性和动态量化

    📝通俗解释：量化就是把高精度数字（32比特浮点数）转换成低精度数字（8比特整数）的技术。分位数、线性、动态是三种不同的量化策略，各有优缺点，适用于不同的场景。

  • 快速的分位数估计：比其他算法快 100 倍

    📝通俗解释：分位数估计是量化过程中的一个重要步骤。bitsandbytes 优化了这个计算过程，比传统方法快100倍，大大减少了模型量化的准备时间。

二、如何才能使用 bitsandbytes？ ​

量化模型的唯一条件是包含 torch.nn.Linear 层，因此量化对于任何模态都可以实现开箱即用。用户可以开箱即用地加载诸如 Whisper（语音识别）、ViT（图像分类）、Blip2（图像问答）之类的 8 比特或 4 比特 (FP4/NF4) 模型。

📝通俗解释：torch.nn.Linear 是 PyTorch 中最基本的线性变换层（就是 y = wx + b），几乎所有的深度学习模型都会用到它。因为这个层在各类模型中都很常见，所以 bitsandbytes 的量化技术可以广泛应用在语音、图像、文字等各种类型的模型上，不需要额外修改模型代码。

三、如何使用 bitsandbytes？ ​

使用 NF4 量化加载 4 比特模型 ​

python
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 配置 NF4 量化参数
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,              # 开启4比特量化加载
    bnb_4bit_quant_type="nf4",      # 使用NF4量化类型
    bnb_4bit_use_double_quant=True, # 开启双重量化，进一步压缩
    bnb_4bit_compute_dtype=torch.bfloat16  # 计算时使用16比特精度
)

# 加载量化后的模型
model_nf4 = AutoModelForCausalLM.from_pretrained(
    model_id, 
    quantization_config=nf4_config
)

📝通俗解释：

  • NF4：一种专门为大语言模型设计的4比特量化方法，比普通的FP4更准确
  • 双重量化（double_quant）：对量化后的参数再进行一次量化，就像把已经压缩的文件再压缩一次，可以进一步节省约10%的显存
  • compute_dtype：虽然模型参数是4比特，但计算时临时用16比特（bfloat16），这样既能节省显存，又能保持一定的计算精度

使用 FP4 量化加载 4 比特模型 ​

python
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 配置 FP4 量化参数
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,              # 开启4比特量化加载
    bnb_4bit_compute_dtype=torch.bfloat16  # 计算时使用16比特精度
)

# 加载量化后的模型
model_fp4 = AutoModelForCausalLM.from_pretrained(
    model_id, 
    quantization_config=quantization_config
)

📝通俗解释：FP4 是另一种4比特量化方法，相比NF4更通用，但在大语言模型上效果略逊于NF4。选择哪种方法取决于具体场景：如果追求更好的模型效果用NF4，如果追求更快的加载速度可以考虑FP4。

总结 ​

bitsandbytes 为大模型的推理和训练提供了高效的显存优化方案：

特性效果
8比特推理显存减半，速度接近无损
4比特量化显存降至1/4，适合消费级GPU运行大模型
8比特优化器训练显存节省75%，可训练更大模型

📝通俗解释：bitsandbytes 就像一个"模型压缩工具箱"，它通过把模型参数从32比特压缩到8比特或4比特，让普通玩家也能在有限显存的情况下运行和训练大模型。比如一个70亿参数的大模型 normally 需要约14GB显存，压缩后只需要3-4GB，一块普通的游戏显卡就能跑起来了。

基于 MIT 许可发布

版权所有 © 2024 LLMs