Views
No views yet
1import torch
2from llmcompressor import oneshot
3from llmcompressor.modifiers.quantization import QuantizationModifier
4from transformers import (
5 AutoModelForCausalLM, # type:ignore[reportPrivateImportUsage]
6 AutoTokenizer, # type:ignore[reportPrivateImportUsage]
7)
8
9MODEL_ID = "tiiuae/falcon-7b"
10
11# Copied from <https://github.com/vllm-project/llm-compressor/blob/9d8a46418f517dd6399e2e9c179805247a7be584/examples/quantization_w8a8_fp8/README.md>
12model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
13tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
14
15# Configure the simple PTQ quantization
16recipe = QuantizationModifier(
17 targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]
18)
19
20# Apply the quantization algorithm.
21oneshot(model=model, recipe=recipe) # type:ignore[arg-type]
22
23# Save the model.
24SAVE_DIR = f"data/models/{MODEL_ID.split('/')[-1]}-FP8-Dynamic"
25model.save_pretrained(SAVE_DIR)
26tokenizer.save_pretrained(SAVE_DIR)