如何在CPU上量化句子转换器模型以在GPU上使用它？

Question

我想使用“Salesforce/SFR-Embedding-Mistral”嵌入模型，但它对于我可以访问的 GPU 分区来说太大了。因此，我考虑量化模型，但找不到可用的预量化版本。

当我尝试使用位和字节对其进行量化时，它尝试将整个模型加载到 GPU 上，这导致了相同的错误。

model = AutoModel.from_pretrained(
    'Salesforce/SFR-Embedding-Mistral',
    trust_remote_code=True,
    device_map='auto',
    torch_dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
)

然后，我尝试先将模型加载到CPU上，然后对其进行量化，然后再将量化后的模型移动到GPU上：

model.to('cpu')
if torch.cuda.is_available():
    model.to('cuda')

但是，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`.

我找到的解决方案，例如这个GitHub问题和这个博客文章，没有帮助或者已经过时了。

Answer 1

消防队。

在这种情况下，您只需加载模型，无需转移到 CUDA。如错误日志中所示，当参数

device_map

为

'auto'

、

device_map='auto'

，并且您使用 4 位或 8 位的量化配置

bnbconfig

时，模型会自动转移到 CUDA。因此，它不能移动两次，并且会出现错误。

因此，仅使用此代码即可工作：

model = AutoModel.from_pretrained(
    'Salesforce/SFR-Embedding-Mistral',
    trust_remote_code=True,
    device_map='auto',
    torch_dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
)

我希望你会发现它很有用:)

亲切的问候，

安德烈斯·埃伦西亚 数据科学和云解决方案

https://www.linkedin.com/in/andres-herencia/