如何在CPU上量化Salesforce/SFR-Embedding-Mistral模型以在GPU使用？

问题概述 #

• 目标使用Salesforce/SFR-Embedding-Mistral嵌入模型，但GPU显存不足无法直接加载
• 直接使用bitsandbytes 4-bit量化配置加载时，模型会先完整加载到GPU触发内存不足错误
• 尝试先加载到CPU量化后迁移至GPU，触发bitsandbytes报错：

  ValueError: `.to` is not supported for `4-bit` or `8-bit` bitsandbytes models. Please use the model as it is, since the model has already been set to the correct devices and cast to the correct `dtype`.
  

可行解决方案 #

核心思路 #

强制模型先在CPU上完成量化，再通过兼容的方式迁移至GPU，避免预加载全量模型到GPU导致OOM，同时绕过.to()方法的限制。

代码实现 #

from transformers import AutoModel, BitsAndBytesConfig
import torch
from accelerate import move_model_to_device

# 定义4-bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# 关键：设置device_map='cpu'，让模型在CPU完成量化，避免预加载到GPU
model = AutoModel.from_pretrained(
    'Salesforce/SFR-Embedding-Mistral',
    trust_remote_code=True,
    device_map='cpu',
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)

# 使用accelerate工具将量化后的模型迁移至GPU（兼容bitsandbytes量化模型）
if torch.cuda.is_available():
    move_model_to_device(model, torch.device('cuda'))

# 验证模型设备
print(next(model.parameters()).device)  # 应输出cuda:0

补充说明 #

1. 为什么device_map='auto'会触发OOM？
   默认情况下，transformers会优先尝试将模型权重预加载到GPU再执行量化操作，导致显存被全量模型占满。设置device_map='cpu'会强制在CPU完成量化，再迁移已压缩的量化权重。

2. 替代迁移方案（无accelerate依赖）
   若无法安装accelerate，可手动遍历模型参数迁移：

   if torch.cuda.is_available():
       # 迁移量化参数至GPU
       for module in model.modules():
           if hasattr(module, 'qweight') or hasattr(module, 'qzeros'):
               module.qweight = module.qweight.to('cuda')
               module.qzeros = module.qzeros.to('cuda')
               if hasattr(module, 'scales'):
                   module.scales = module.scales.to('cuda')
       # 设置模型计算设备
       model = model.to('cuda', dtype=torch.bfloat16)
   

3. 进一步显存优化

   • 关闭bnb_4bit_use_double_quant可减少少量显存占用（但会损失一点精度）
   • 若仍显存不足，可改用8-bit量化（将load_in_4bit=True改为load_in_8bit=True），兼容性更好但显存占用略高

内容的提问来源于stack exchange，提问作者Firevince