Transformers 文档

FBGEMM FP8

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

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")
< > 更新 在 GitHub 上