前言

量化(Quantization)是模型性能优化技术中的一种，在压缩内存的同时可以实现更快速的计算。

第三方库

pip install bitsandbytes

Step 1 量化模型并保存

关键代码

model_name_or_path = "codellama/CodeLlama-34b-Instruct-hf"
tokenizer_name_or_path = model_name_or_path

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path, 
    local_files_only=True,# trust_remote_code=True,
    quantization_config = BitsAndBytesConfig(
		# 量化数据类型设置
        bnb_4bit_quant_type="nf4",
		# 量化数据的数据格式
        bnb_4bit_compute_dtype=torch.bfloat16
    ),
    device_map='auto', 
    torch_dtype=torch.bfloat16,
)


output = "codellama/CodeLlama-34b-Instruct-hf-4bit"
if not os.path.exists(output):
    os.mkdir(output)

model.save_pretrained(output)
print("done")
exit(0)

Step 2 加载量化后的模型

关键代码

model_name_or_path = "codellama/CodeLlama-34b-Instruct-hf-4bit"
tokenizer_name_or_path = model_name_or_path

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path, 
    local_files_only=True,
    device_map='auto', 
    torch_dtype=torch.bfloat16,
    load_in_4bit=True, # 这句是关键
)

# 其余的使用和量化前没区别

性能提升

推理速度接近量化前的两倍。

占用最高显存约为量化前的二分之一。(这意味着batch的大小可以翻一倍？)

这是生成时最高的显存占用情况

附：量化前后的本地文件

会发现大小没什么变化，并且最后四个文件(special_tokens_map.json,tokenizer_config.json,tokenizer.json,tokenizer.model) 不会生成，需要从原来的目录下复制到新的目录下

codellama/CodeLlama-34b-Instruct-hf

codellama/CodeLlama-34b-Instruct-hf-4bit 

参考网址

使用 Transformers 量化 Meta AI LLaMA2 中文版大模型

Quantize 🤗 Transformers models