Transformers 文档

FBGEMM FP8

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")
< > Update on GitHub