FBGEMM FP8
使用FBGEMM FP8量化方法,您可以将模型量化为FP8(W8A8):
- 权重将按通道量化为8位(FP8)
- 激活将按每个令牌量化为8位(FP8)
它依赖于FBGEMM库,该库为小批量大小提供了高效的低精度通用矩阵乘法,并支持减少精度损失的技术,如行量化(row-wise quantization)和异常值感知量化(outlier-aware quantization)。
你需要一个计算能力>=9的GPU(例如H100)
在开始之前,请确保以下库已安装并更新到最新版本:
pip install --upgrade accelerate fbgemm-gpu torch
如果您在使用fbgemm-gpu和torch库时遇到问题,您可能需要安装nightly版本。您可以按照这里的说明进行操作。
默认情况下,权重以全精度(torch.float32)加载,无论权重实际存储的数据类型是什么,例如torch.float16。设置torch_dtype="auto"
以加载模型config.json
文件中定义的数据类型,以自动加载内存最优的数据类型。
from transformers import FbgemmFp8Config, AutoModelForCausalLM, AutoTokenizer
model_name = "meta-llama/Meta-Llama-3-8B"
quantization_config = FbgemmFp8Config()
quantized_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto", quantization_config=quantization_config)
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
output = quantized_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))
量化模型可以通过“saved_pretrained”保存,并通过“from_pretrained”再次使用。
quant_path = "/path/to/save/quantized/model"
model.save_pretrained(quant_path)
model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto")