细粒度 FP8

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

权重针对每个 2D 块量化为 8 位 (weight_block_size=(128, 128))。
激活针对每个令牌的每个组量化为 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 上更新

Transformers

细粒度 FP8