显存不够又想用某个模型时的模型量化操作

鱼鱼9901

已于 2024-01-26 22:41:32 修改

阅读量524

点赞数 3

分类专栏： nlp 文章标签：人工智能 linux 深度学习

于 2024-01-26 21:37:55 首次发布

本文链接：https://blog.csdn.net/weixin_72100405/article/details/135875328

版权

nlp 专栏收录该内容

15 篇文章

订阅专栏

本文介绍了如何使用HuggingFace的Transformers库中的BitsAndBytesConfig对预训练的AutoModelForCausalLM模型进行8位量化，并在考虑到GPU内存限制的情况下，进行了设备映射和代码信任设置。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

from transformers import AutoTokenizer, AutoModelForCausalLM,BitsAndBytesConfig
quantization_config= BitsAndBytesConfig(load_in_8bit=True)

tokenizer = AutoTokenizer.from_pretrained(path,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(path,
device_map="cuda:0",
trust_remote_code=True,
quantization_config=quantization_config,          max_memory=torch.cuda.get_device_properties(0).total_memory
).eval()