|
```python |
|
import time |
|
from datasets import load_dataset |
|
from transformers import AutoTokenizer |
|
from rich import print |
|
from llmcompressor.modifiers.quantization import GPTQModifier |
|
from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot |
|
|
|
# Select model and load it. |
|
MODEL_ID = "Unbabel/TowerInstruct-7B-v0.1" |
|
model = SparseAutoModelForCausalLM.from_pretrained( |
|
MODEL_ID, |
|
device_map="auto", |
|
torch_dtype="auto", |
|
) |
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
|
|
|
# Select calibration dataset. |
|
DATASET_ID = "neuralmagic/LLM_compression_calibration" |
|
DATASET_SPLIT = "train" |
|
|
|
# Select number of samples. 512 samples is a good place to start. |
|
# Increasing the number of samples can improve accuracy. |
|
NUM_CALIBRATION_SAMPLES = 756 |
|
MAX_SEQUENCE_LENGTH = 2048 |
|
|
|
# Load dataset and preprocess. |
|
ds = load_dataset(DATASET_ID, split=DATASET_SPLIT) |
|
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES)) |
|
|
|
|
|
def preprocess(example): |
|
return { |
|
"text": tokenizer.apply_chat_template( |
|
example["messages"], |
|
tokenize=False, |
|
) |
|
} |
|
|
|
|
|
ds = ds.map(preprocess) |
|
|
|
|
|
# Tokenize inputs. |
|
def tokenize(sample): |
|
return tokenizer( |
|
sample["text"], |
|
padding=False, |
|
max_length=MAX_SEQUENCE_LENGTH, |
|
truncation=True, |
|
add_special_tokens=False, |
|
) |
|
|
|
|
|
ds = ds.map(tokenize, remove_columns=ds.column_names) |
|
|
|
# Configure the quantization algorithm to run. |
|
# * quantize the weights to 4 bit with GPTQ with a group size 128 |
|
# Note: to reduce GPU memory use `sequential_update=False` |
|
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) |
|
print(recipe) |
|
|
|
# Apply algorithms. |
|
oneshot( |
|
model=model, |
|
dataset=ds, |
|
recipe=recipe, |
|
max_seq_length=MAX_SEQUENCE_LENGTH, |
|
num_calibration_samples=NUM_CALIBRATION_SAMPLES, |
|
) |
|
|
|
# Confirm generations of the quantized model look sane and measure generation time. |
|
print("\n\n") |
|
print("========== SAMPLE GENERATION ==============") |
|
input_text = "Translate the following text from Portuguese into English.\nPortuguese: Um grupo de investigadores lançou um novo modelo para tarefas relacionadas com tradução.\nEnglish:" |
|
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda") |
|
start_time = time.time() |
|
output = model.generate(input_ids, max_new_tokens=100) |
|
end_time = time.time() |
|
generation_time = end_time - start_time |
|
print(tokenizer.decode(output[0])) |
|
print(f"Generation time: {generation_time:.2f} seconds") |
|
print("==========================================\n\n") |
|
|
|
# Save to disk compressed. |
|
SAVE_DIR = MODEL_ID.split("/")[1] + "-quantized.w4a16" |
|
model.save_pretrained(SAVE_DIR, save_compressed=True) |
|
tokenizer.save_pretrained(SAVE_DIR) |
|
``` |