FBGEMM FP8
使用 FBGEMM FP8 量化方法,您可以将模型量化为 FP8 (W8A8)
- 权重将以每通道 8 位(FP8)进行量化
- 激活将以每个令牌 8 位(FP8)进行量化
它依赖于 FBGEMM 库,该库为小批量大小提供高效的低精度通用矩阵乘法,并支持行级量化和异常值感知量化等精度损失最小化技术。
您需要一个计算能力 >= 9 的 GPU(例如 H100)
在开始之前,请确保以下库已安装其最新版本
pip install --upgrade accelerate fbgemm-gpu torch
如果您在使用 fbgemm-gpu 和 torch 库时遇到问题,您可能需要安装 nightly 版本。您可以按照 此处 的说明操作。
from transformers import FbgemmFp8Config, AutoModelForCausalLM, AutoTokenizer
model_name = "meta-llama/Meta-Llama-3-8B"
quantization_config = FbgemmFp8Config()
quantized_model = AutoModelForCausalLM.from_pretrained(model_name, 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")