Compressing Massive Language Fashions (LLMs) | by Shaw Talebi | Aug, 2024

We begin by importing a number of useful libraries.

from datasets import load_dataset

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import DistilBertForSequenceClassification, DistilBertConfig

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.knowledge import DataLoader

from sklearn.metrics import accuracy_score, precision_recall_fscore_support

Then, we load our dataset from the Hugging Face Hub. This contains coaching (2100 rows), testing (450 rows), and validation (450 rows) units.

knowledge = load_dataset("shawhin/phishing-site-classification")

Subsequent, we load in our instructor mannequin. To assist velocity up coaching, I loaded the mannequin onto a T4 GPU that was freely obtainable on Google Colab.

# use Nvidia GPU
machine = torch.machine('cuda')

# Load instructor mannequin and tokenizer
model_path = "shawhin/bert-phishing-classifier_teacher"

tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)
.to(machine)

The instructor mannequin is a fine-tuned model of Goolge’s bert-base-uncased that performs binary classification on phishing web site URLs. The code to coach the instructor mannequin is out there on GitHub.

For the coed mannequin, we initialize a brand new mannequin from scratch primarily based on distilbert-base-uncased. We modify the structure by eradicating two layers and 4 consideration heads from the remaining layers.

# Load scholar mannequin
my_config = DistilBertConfig(n_heads=8, n_layers=4) # drop 4 heads per layer and a pair of layers

student_model = DistilBertForSequenceClassification
.from_pretrained("distilbert-base-uncased",
config=my_config,)
.to(machine)

Earlier than we will practice our scholar mannequin, we might want to tokenize the dataset. That is necessary as a result of the fashions anticipate enter textual content to be represented in a selected manner.

Right here, I pad examples primarily based on every batch’s longest instance. This permits the batches to be represented as a PyTorch tensor.

# outline textual content preprocessing
def preprocess_function(examples):
return tokenizer(examples["text"], padding='max_length', truncation=True)

# tokenize all datasetse
tokenized_data = knowledge.map(preprocess_function, batched=True)
tokenized_data.set_format(sort='torch',
columns=['input_ids', 'attention_mask', 'labels'])

One other necessary step earlier than coaching is defining an analysis technique for our fashions throughout coaching. Beneath, I outline a perform that computes the accuracy, precision, recall, and F1 rating given a mannequin and dataset.

# Operate to judge mannequin efficiency
def evaluate_model(mannequin, dataloader, machine):
mannequin.eval() # Set mannequin to analysis mode
all_preds = []
all_labels = []

# Disable gradient calculations
with torch.no_grad():
for batch in dataloader:
input_ids = batch['input_ids'].to(machine)
attention_mask = batch['attention_mask'].to(machine)
labels = batch['labels'].to(machine)

# Ahead go to get logits
outputs = mannequin(input_ids, attention_mask=attention_mask)
logits = outputs.logits

# Get predictions
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.lengthen(preds)
all_labels.lengthen(labels.cpu().numpy())

# Calculate analysis metrics
accuracy = accuracy_score(all_labels, all_preds)
precision, recall, f1, _ = precision_recall_fscore_support(all_labels,
all_preds,
common='binary')

return accuracy, precision, recall, f1

Now, we’re prepared to start the coaching course of. To permit our scholar mannequin to study from each the bottom fact labels within the coaching set (i.e., arduous targets) and the instructor mannequin’s logits (i.e., delicate targets), we should assemble a particular loss perform that considers each targets.

That is accomplished by combining the KL divergence of the coed and instructor’s output chance distribution with the cross entropy loss of the coed’s logits with the bottom fact.

# Operate to compute distillation and hard-label loss
def distillation_loss(student_logits, teacher_logits,
true_labels, temperature, alpha):
# Compute delicate targets from instructor logits
soft_targets = nn.practical.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.practical.log_softmax(student_logits / temperature, dim=1)

# KL Divergence loss for distillation
distill_loss = nn.practical.kl_div(student_soft,
soft_targets,
discount='batchmean') * (temperature ** 2)

# Cross-entropy loss for arduous labels
hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)

# Mix losses
loss = alpha * distill_loss + (1.0 - alpha) * hard_loss

return loss

Subsequent, we outline our hyperparameters, optimizer, and practice/take a look at datasets.

# hyperparameters
batch_size = 32
lr = 1e-4 #5e-5
num_epochs = 5
temperature = 2.0
alpha = 0.5

# outline optimizer
optimizer = optim.Adam(student_model.parameters(), lr=lr)

# create coaching knowledge loader
dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)
# create testing knowledge loader
test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)

Lastly, we practice our scholar mannequin utilizing PyTorch.

# put scholar mannequin in practice mode
student_model.practice()

# practice mannequin
for epoch in vary(num_epochs):
for batch in dataloader:
# Put together inputs
input_ids = batch['input_ids'].to(machine)
attention_mask = batch['attention_mask'].to(machine)
labels = batch['labels'].to(machine)

# Disable gradient calculation for instructor mannequin
with torch.no_grad():
teacher_outputs = teacher_model(input_ids,
attention_mask=attention_mask)
teacher_logits = teacher_outputs.logits

# Ahead go by way of the coed mannequin
student_outputs = student_model(input_ids,
attention_mask=attention_mask)
student_logits = student_outputs.logits

# Compute the distillation loss
loss = distillation_loss(student_logits, teacher_logits, labels,
temperature, alpha)

# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()

print(f"Epoch {epoch + 1} accomplished with loss: {loss.merchandise()}")

# Consider the instructor mannequin
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =
evaluate_model(teacher_model, test_dataloader, machine)

print(f"Trainer (take a look at) - Accuracy: {teacher_accuracy:.4f},
Precision: {teacher_precision:.4f},
Recall: {teacher_recall:.4f},
F1 Rating: {teacher_f1:.4f}")

# Consider the coed mannequin
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, test_dataloader, machine)

print(f"Pupil (take a look at) - Accuracy: {student_accuracy:.4f},
Precision: {student_precision:.4f},
Recall: {student_recall:.4f},
F1 Rating: {student_f1:.4f}")
print("n")

# put scholar mannequin again into practice mode
student_model.practice()

The coaching outcomes are proven within the screenshot under. Remarkably, by the tip of coaching, the scholar mannequin outperformed the instructor throughout all analysis metrics!

Data distillation coaching outcomes. Picture by creator.

As a last step, we will consider the fashions on the impartial validation set, i.e., knowledge not utilized in coaching mannequin parameters or tuning hyperparameters.

# create testing knowledge loader
validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)

# Consider the instructor mannequin
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =
evaluate_model(teacher_model, validation_dataloader, machine)
print(f"Trainer (validation) - Accuracy: {teacher_accuracy:.4f},
Precision: {teacher_precision:.4f},
Recall: {teacher_recall:.4f},
F1 Rating: {teacher_f1:.4f}")

# Consider the coed mannequin
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, validation_dataloader, machine)
print(f"Pupil (validation) - Accuracy: {student_accuracy:.4f},
Precision: {student_precision:.4f},
Recall: {student_recall:.4f},
F1 Rating: {student_f1:.4f}")

Right here, once more, we see the coed outperform the instructor.

Mannequin performances on the validation set. Picture by creator.

Up to now, we’ve decreased our mannequin from 109M parameters (438 MB) to 52.8M parameters (211 MB). Nevertheless, we will go one step additional and quantize the coed mannequin.

First, we push the mannequin of the Hugging Face Hub.

student_model.push_to_hub("shawhin/bert-phishing-classifier_student")

Then, we will load it again in utilizing 4-bit quantization. For that, we will use the BitsAndBytes integration within the transformers library.

We arrange the config to retailer mannequin parameters utilizing the 4-bit NormalFloat knowledge sort described within the QLoRA paper and the bfloat16 for computation [10].

from transformers import BitsAndBytesConfig

# load mannequin in mannequin as 4-bit
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype = torch.bfloat16,
bnb_4bit_use_double_quant=True
)

model_nf4 = AutoModelForSequenceClassification.from_pretrained(model_id,
device_map=machine,
quantization_config=nf4_config)

We are able to then consider our quantized mannequin on the validation set.

# Consider the coed mannequin
quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 =
evaluate_model(model_nf4, validation_dataloader, machine)

print("Put up-quantization Efficiency")
print(f"Accuracy: {quantized_accuracy:.4f},
Precision: {quantized_precision:.4f},
Recall: {quantized_recall:.4f},
F1 Rating: {quantized_f1:.4f}")

Pupil mannequin efficiency on the validation set after quantization. Picture by creator.

As soon as once more, we see a small efficiency enchancment after compression. An intuitive clarification for that is Occam’s Razor precept, which states that easier fashions are higher.

On this case, the mannequin could also be overparameterized for this binary classification activity. Thus, simplifying the mannequin leads to higher efficiency.

Whereas trendy giant language fashions (LLMs) exhibit spectacular efficiency on varied duties, their scale presents challenges in deploying them in real-world settings.

Current improvements in mannequin compression methods assist mitigate these challenges by lowering the computational value of LLM options. Right here, we mentioned three broad classes of compression methods (Quantization, Pruning, and Data Distillation) and walked by way of an instance implementation in Python.

Extra on LLMs 👇

Shaw Talebi

Massive Language Fashions (LLMs)

My web site: https://www.shawhintalebi.com/