關注TechLead，復旦AI博士，分享AI領域全維度知識與研究。擁有10+年AI領域研究經驗、復旦機器人智慧實驗室成員，國家級大學生賽事評審專家，發表多篇SCI核心期刊學術論文，上億營收AI產品研發負責人。
如何在不犧牲效能的情況下將大型語言模型縮小十倍
雖然LLM的巨大規模賦予了它們在各種用例中的出色效能，但這也在其應用於現實世界問題時帶來了挑戰。在本文中，我將討論如何透過壓縮LLM來克服這些挑戰。我將從概述關鍵概念開始，接著透過Python程式碼展示一個具體的示例。
2023年AI領域的口號是"越大越好"，提升語言模型的公式非常簡單：更多的資料 + 更多的引數 + 更多的計算資源 = 更好的效能。
雖然這仍然是目前的趨勢，但處理1000億以上引數的模型顯然存在挑戰。例如，一個具有1000億引數的模型僅在FP16格式下儲存就需要200GB的空間！
不用說，大多數消費裝置（如手機、平板電腦、膝上型電腦）無法處理如此龐大的模型。但……如果我們可以讓模型變小呢？
模型壓縮
模型壓縮旨在在不犧牲效能的前提下減少機器學習模型的大小。對於（大型）神經網路，這可行，因為它們通常是過引數化的（即由冗餘的計算單元組成）。
模型壓縮的主要好處是降低推理成本。這意味著功能強大的ML模型可以更廣泛地被使用（例如在本地膝上型電腦上執行LLM），並且能夠以更低的成本將AI整合到消費產品中，還支援裝置上的推理，從而保護使用者隱私和安全。
三種壓縮模型的方法
模型壓縮有多種技術。這裡我將重點介紹三種廣泛使用的類別。
- 量化——使用更低精度的資料型別表示模型
- 剪枝——從模型中刪除不必要的元件
- 知識蒸餾——透過較大的模型訓練較小的模型
注意：這些方法是相互獨立的。因此，可以組合來自多個類別的技術，以實現最大的壓縮效果！
量化
雖然"量化"聽起來像是一個複雜的術語，但它實際上是一個簡單的概念。它包括降低模型引數的精度。你可以把它想象成將高解析度影像轉換為低解析度影像，同時保持圖片的核心屬性。
兩類常見的量化技術是後訓練量化（PTQ）和量化感知訓練（QAT）。
後訓練量化（PTQ）
對於一個神經網路，後訓練量化（PTQ）透過將引數替換為低精度的資料型別（例如從FP16轉換為INT-8）來壓縮模型。這是減少模型計算需求最快且最簡單的方法之一，因為它不需要額外的訓練或資料標註。
雖然這是降低模型成本的相對簡單的方法，但以這種方式過度量化（例如從FP16到INT4）通常會導致效能下降，這限制了PTQ的潛在收益。
量化感知訓練（QAT）
對於需要更高壓縮的情況，可以透過從零開始用低精度資料型別訓練模型來克服PTQ的侷限性。這就是量化感知訓練（QAT）的思想。
雖然這種方法在技術上更具挑戰性，但它可以生成一個顯著更小且表現良好的模型。例如，BitNet架構使用了一種三進位制資料型別（即1.58位），匹配了原始Llama模型的效能！
當然，PTQ和從零開始的QAT之間存在較大的技術鴻溝。一種介於兩者之間的方法是量化感知微調，它包括在量化後對預訓練模型進行額外的訓練。
剪枝
剪枝的目的是刪除對模型效能影響較小的元件。這很有效，因為ML模型（尤其是大型模型）往往會學習到冗餘和噪聲結構。
一個比喻是修剪樹木中的枯枝。去除它們可以減少樹的大小而不會傷害樹。
剪枝方法可以分為兩類：非結構化剪枝和結構化剪枝。
非結構化剪枝
非結構化剪枝是從神經網路中移除不重要的權重（即將它們設定為零）。例如，早期的工作如Optimal Brain Damage和Optimal Brain Surgeon透過估計每個引數對損失函式的影響來計算其重要性分數。
最近，基於幅值的方法（即移除絕對值最小的權重）變得更受歡迎，因為它們簡單且易於擴充套件。
雖然非結構化剪枝可以顯著減少引數數量，但這些收益需要特殊硬體才能實現。非結構化剪枝會導致稀疏矩陣操作（即乘以大量零的矩陣），而標準硬體無法比非稀疏操作更有效地執行這些操作。
結構化剪枝
相對而言，結構化剪枝是從神經網路中刪除整個結構（例如注意力頭、神經元和層）。這樣可以避免稀疏矩陣操作的問題，因為整個矩陣可以從模型中刪除，而不是個別引數。
雖然有多種方法可以確定要剪枝的結構，但原則上，它們都試圖刪除對效能影響最小的結構。
知識蒸餾
知識蒸餾是將知識從一個（較大的）教師模型傳遞到一個（較小的）學生模型。一種方法是透過教師模型生成預測，並使用這些預測來訓練學生模型。透過學習教師模型的輸出logits（即所有可能的下一個標記的機率），學生模型獲得了比原始訓練資料更豐富的資訊，從而提高了效能。
最近的蒸餾應用完全摒棄了logits的需求，而是透過教師模型生成的合成資料進行學習。一個著名的例子是史丹佛的Alpaca模型，它使用來自OpenAI的text-davinci-003（即原始ChatGPT模型）的合成資料微調了LLaMa 7B（基礎）模型，使其能夠遵循使用者指令。
示例程式碼：透過知識蒸餾和量化壓縮文字分類器
在基本瞭解了各種壓縮技術後，讓我們看一個如何在Python中進行壓縮的實際示例。這裡，我們將壓縮一個有1億引數的模型，該模型用於分類URL是否安全（即釣魚網站）。
我們首先使用知識蒸餾將1億引數的模型壓縮到5000萬引數。然後，使用4位量化進一步將記憶體佔用減少3倍，得到的最終模型比原始模型小7倍。
示例程式碼可以在 https://github.com/ShawhinT/YouTube-Blog/tree/main/LLMs/model-compression 上找到。教師模型 https://huggingface.co/shawhin/bert-phishing-classifier_teacher 、學生模型https://huggingface.co/shawhin/bert-phishing-classifier_student、 4位學生模型 https://huggingface.co/shawhin/bert-phishing-classifier_student_4bit 、資料集https://huggingface.co/datasets/shawhin/phishing-site-classification 可以在Hugging Face Hub上免費獲取。
首先，我們匯入一些有用的庫。
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.data import DataLoader
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
然後，我們從Hugging Face Hub載入資料集。包括訓練集（2100行）、測試集（450行）和驗證集（450行）。
data = load_dataset("shawhin/phishing-site-classification")
接下來，我們載入教師模型。為了加速訓練，我將模型載入到Google Colab上提供的免費T4 GPU上。
device = torch.device('cuda')
model_path = "shawhin/bert-phishing-classifier_teacher"
tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pre
trained(model_path).to(device)
教師模型是Google的 bert-base-uncased https://huggingface.co/google-bert/bert-base-uncased 的微調版本，執行對釣魚網站URL的二分類。訓練教師模型的程式碼可在 GitHub https://github.com/ShawhinT/YouTube-Blog/tree/main/LLMs/model-compression 上找到。
對於學生模型，我們從 distilbert-base-uncased https://huggingface.co/distilbert/distilbert-base-uncased 初始化一個新模型。我們透過移除兩層和剩餘層的四個注意力頭修改了架構。
my_config = DistilBertConfig(n_heads=8, n_layers=4)
student_model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", config=my_config).to(device)
在訓練學生模型之前，我們需要對資料集進行標記。這是必要的，因為模型期望輸入文字以特定的方式表示。
在這裡，我根據每個批次的最長示例填充樣本。這使批次能夠表示為PyTorch張量。
def preprocess_function(examples):
return tokenizer(examples["text"], padding='max_length', truncation=True)
tokenized_data = data.map(preprocess_function, batched=True)
tokenized_data.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
訓練前的另一個重要步驟是定義一個評估策略，用於在訓練期間評估我們的模型。在下面，我定義了一個函式，該函式在給定模型和資料集的情況下計算準確率、精確率、召回率和F1得分。
def evaluate_model(model, dataloader, device):
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
accuracy = accuracy_score(all_labels, all_preds)
precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')
return accuracy, precision, recall, f1
現在我們可以開始訓練了。為了讓學生模型同時學習訓練集中的真實標籤（即硬目標）和教師模型的logits（即軟目標），我們需要構建一個特殊的損失函式，該函式考慮到兩種目標。
這可以透過結合學生和教師輸出機率分佈的KL散度與學生logits與真實標籤的交叉熵損失來實現。
def distillation_loss(student_logits, teacher_logits, true_labels, temperature, alpha):
soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)
distill_loss = nn.functional.kl_div(student_soft, soft_targets, reduction='batchmean') * (temperature ** 2)
hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)
loss = alpha * distill_loss + (1.0 - alpha) * hard_loss
return loss
接下來，我們定義超引數、最佳化器和訓練/測試資料集。
batch_size = 32
lr = 1e-4
num_epochs = 5
temperature = 2.0
alpha = 0.5
optimizer = optim.Adam(student_model.parameters(), lr=lr)
dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)
test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)
最後，我們使用PyTorch訓練學生模型。
student_model.train()
for epoch in range(num_epochs):
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
with torch.no_grad():
teacher_outputs = teacher_model(input_ids, attention_mask=attention_mask)
teacher_logits = teacher_outputs.logits
student_outputs = student_model(input_ids, attention_mask=attention_mask)
student_logits = student_outputs.logits
loss = distillation_loss(student_logits, teacher_logits, labels, temperature, alpha)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1} completed with loss: {loss.item()}")
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 = evaluate_model(teacher_model, test_dataloader, device)
print(f"Teacher (test) - Accuracy: {teacher_accuracy:.4f}, Precision: {teacher_precision:.4f}, Recall: {teacher_recall:.4f}, F1 Score: {teacher_f1:.4f}")
student_accuracy, student_precision, student_recall, student_f1 = evaluate_model(student_model, test_dataloader, device)
print(f"Student (test) - Accuracy: {student_accuracy:.4f}, Precision: {student_precision:.4f}, Recall: {student_recall:.4f}, F1 Score: {student_f1:.4f}")
print("\n")
student_model.train()
訓練結果如下圖所示。令人驚訝的是，訓練結束時，學生模型在所有評估指標上都超過了教師模型！
接下來，我們可以在獨立的驗證集上評估模型，即未用於訓練模型引數或調整超引數的資料。
validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 = evaluate_model(teacher_model, validation_dataloader, device)
print(f"Teacher (validation) - Accuracy: {teacher_accuracy:.4f}, Precision: {teacher_precision:.4f}, Recall: {teacher_recall:.4f}, F1 Score: {teacher_f1:.4f}")
student_accuracy, student_precision, student_recall, student_f1 = evaluate_model(student_model, validation_dataloader, device)
print(f"Student (validation) - Accuracy: {student_accuracy:.4f}, Precision: {student_precision:.4f}, Recall: {student_recall:.4f}, F1 Score: {student_f1:.4f}")
在這裡，我們再次看到學生模型優於教師模型。
到目前為止，我們已經將模型從1.09億引數（438 MB）縮小到5280萬引數（211 MB）。然而，我們可以更進一步，對學生模型進行量化。
首先，我們將模型推送到 Hugging Face Hub https://huggingface.co/shawhin/bert-phishing-classifier_student 。
student_model.push_to_hub("shawhin/bert-phishing-classifier_student")
然後，我們可以使用4位量化載入模型。為此，我們可以使用transformers庫中的BitsAndBytes整合。
我們設定配置以使用 QLoRA https://medium.com/towards-data-science/qlora-how-to-fine-tune-an-llm-on-a-single-gpu-4e44d6b5be32 論文中描述的4位NormalFloat資料型別儲存模型引數，並使用bfloat16進行計算。
from transformers import BitsAndBytesConfig
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=device, quantization_config=nf4_config)
接下來，我們可以在驗證集上評估量化模型。
quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 = evaluate_model(model_nf4, validation_dataloader, device)
print("Post-quantization Performance")
print(f"Accuracy: {quantized_accuracy:.4f}, Precision: {quantized_precision:.4f}, Recall: {quantized_recall:.4f}, F1 Score: {quantized_f1:.4f}")
我們再次看到壓縮後的效能略有提升（縮小到62.7MB）。一個直觀的解釋是奧卡姆剃刀原理，即越簡單的模型越好。
在這種情況下，模型可能對於這個二分類任務來說過度引數化了。因此，簡化模型帶來了更好的效能。
原文連結：https://medium.com/towards-data-science/compressing-large-language-models-llms-9f406eea5b5e
本文由部落格一文多發平臺 OpenWrite 釋出！