In [ ]:
pip install --upgrade accelerate fbgemm-gpu torch
如果安装 fbgemm-gpu 和 torch 库时遇到问题，可能需要安装夜间版。你可以按照 这里的说明 进行操作。
默认情况下，无论权重实际存储的数据类型（如 torch.float16），权重都会以全精度（torch.float32）加载。设置 torch_dtype="auto" 可以根据模型的 config.json 文件中定义的数据类型自动加载最节省内存的数据类型。
In [ ]:
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 = "我们今晚吃什么？"
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 重新加载使用。
In [ ]:
quant_path = "/path/to/save/quantized/model"
model.save_pretrained(quant_path)
model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto")