Transformers 文档

细粒度 FP8

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

细粒度 FP8

细粒度 FP8 量化将权重和激活量化为 fp8。

  • 权重对每个 2D 块(`weight_block_size=(128, 128)`)量化为 8 位。
  • 激活对每个令牌的每个组进行 8 位量化。组值与输入通道中的权重匹配(默认为 128)。

FP8 量化支持 DeepSeek-V3 和 DeepSeek-R1。

您需要一个计算能力 >=9 (H100) 的 GPU,并安装与您的 GPU 的 CUDA 版本兼容的 PyTorch 版本。

安装 Accelerate 并升级到最新版本的 PyTorch。

pip install --upgrade accelerate torch

创建一个 FineGrainedFP8Config 类,并将其传递给 from_pretrained() 进行量化。默认情况下,无论权重以何种实际数据类型存储,权重都以全精度(`torch.float32`)加载。设置 `torch_dtype="auto"` 以在模型 `config.json` 文件中定义的数据类型中加载权重,从而自动加载最节省内存的数据类型。

from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Meta-Llama-3-8B"
quantization_config = FineGrainedFP8Config()
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))

使用 save_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 上更新