大语言模型的压缩技术

尽管人们对越来越大的语言模型一直很感兴趣，但MistralAI 向我们表明，规模只是相对而言的，而对边缘计算日益增长的兴趣促使我们使用小型语言获得不错的结果。压缩技术提供了一种替代方法。在本文中，我将解释这些技术，并提供一些简单的代码片段作为示例。

模型压缩是在不影响机器学习模型有效性的情况下最小化其大小的行为。由于大型神经网络经常因过度参数化而包含冗余计算单元，因此这种方法对它们非常有效。

压缩意味着减少参数数量或整体内存占用，从而减小模型大小（例如从 10 GB 到 9 GB）。此过程有助于提高模型在存储和推理速度方面的效率，使其更容易在资源有限的环境中部署。常见的模型压缩技术包括：

量化：通过改变模型权重的精度（例如从 32 位浮点数到 8 位整数）来减少内存占用。
修剪：删除不太重要的权重或神经元，减少参数的数量。
知识提炼：训练较小的模型（学生）来模仿较大模型（老师）的行为，将知识提炼为具有类似性能的压缩版本。
权重共享：通过设计或后期训练，在不同层之间使用共享权重来减少存储要求。

1. 模型量化

模型量化通过将权重或激活的精度表示（通常为 32 位或 16 位）转换为较低精度表示（例如 8 位、4 位甚至二进制）来压缩 LLM。我们可以量化权重、激活函数或使用其他技巧：

权重量化：神经网络使用的权重通常存储为 32 位或 16 位浮点数。量化将这些权重降低到较低的位宽，例如 8 位整数 (INT8) 或 4 位整数 (INT4)。这是通过将原始权重范围映射到具有较少位的较小范围来实现的，从而显著减少内存使用量。

激活量化：与权重类似，激活（推理期间的层输出）可以量化为较低的精度。通过用更少的位表示激活，模型在推理期间的内存占用量会减少。

量化感知训练 (QAT)：在 QAT 中，模型在模拟量化的同时进行训练，使其能够适应较低的精度。这有助于保持准确性，因为模型学会了对量化效应更加稳健（查看Tailor 等人的Arxiv）。

训练后量化 (PTQ)：此方法涉及以全精度正常训练模型，然后应用量化。虽然 PTQ 更简单、更快速，但与 QAT 相比，它可能会导致准确率大幅下降（如Wang 等人在 NIPS2021中所述）。

使用 bitsandbytes 可以非常轻松地实现权重量化。安装库：

pip install torch transformers bitsandbytes

例如，对于 GPT2，运行代码：


from transformers import AutoModelForCausalLM, AutoTokenizer
import torch# Specify the model you want to use
model_name = "gpt2"  # You can replace this with any other LLM model
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Load the model with 8-bit quantization using bitsandbytes
model = AutoModelForCausalLM.from_pretrained(model_name,load_in_8bit=True,  # Enable 8-bit quantizationdevice_map="auto"   # Automatically allocate to available device (CPU/GPU)
)
# Example text for inference
input_text = "Weight Quantization is an efficient technique for compressing language models."
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
# Generate text
with torch.no_grad():output_ids = model.generate(input_ids, max_length=50)
# Decode and print the generated text
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

2. 修剪

修剪会移除不必要或不太重要的权重、神经元或整个层，就像从树上移除不必要的树枝一样。这可以减小模型的大小、加快推理速度并降低内存和计算要求，使其更高效，同时尽可能保持原始性能。

这不像量化那么简单，因为我们首先需要找到冗余的东西。例如，我们需要找到冗余参数，然后在没有这些参数的情况下对模型进行微调。

最常见的情况是，我们会移除权重、神经元或层，但人们对注意力头修剪（特定于基于 Transformer 的模型）的兴趣日益浓厚，将其作为一种结构化修剪形式（查看Wang 等人的Arxiv）。在这里，每个注意力层都有多个头。有些头对模型性能的贡献比其他头更大，因此注意力头修剪会移除不太重要的头。

修剪的示例代码如下，我们从 GPT2 模型中删除一定比例的权重：


import torch
import torch.nn.utils.prune as prune
from transformers import AutoModelForCausalLM, AutoTokenizer# Load the pretrained model and tokenizer
model_name = "gpt2"  # You can replace this with any other LLM model
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Define a pruning method (here we use L1 unstructured pruning)
def prune_model_layer(layer, amount=0.3):# Prune 30% of the weights with the lowest L1 norm in the linear layersfor name, module in layer.named_modules():if isinstance(module, torch.nn.Linear):prune.l1_unstructured(module, name="weight", amount=amount)print(f"Pruned layer {name} with amount {amount}")
# Apply pruning to all transformer layers in the model
for layer in model.transformer.h:prune_model_layer(layer, amount=0.3)  # Prune 30% of the weights
# Check the sparsity of the model
total_params = 0
pruned_params = 0
for name, module in model.named_modules():if isinstance(module, torch.nn.Linear):total_params += module.weight.nelement()pruned_params += torch.sum(module.weight == 0).item()
print(f"Total parameters: {total_params}")
print(f"Pruned parameters: {pruned_params}")
print(f"Sparsity: {pruned_params / total_params:.2%}")
# Test the pruned model on a sample input
input_text = "Pruning is an effective way to compress language models."
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
# Generate text using the pruned model
with torch.no_grad():output_ids = model.generate(input_ids, max_length=50)
# Decode and print the generated text
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

3. 模型蒸馏

模型蒸馏是一种将“知识”从一个大型、更复杂的模型（称为教师模型）转移到一个较小、更简单的模型（称为学生模型）的技术，该模型可能具有较少的参数。这一过程使学生模型能够达到接近教师模型的性能，同时在规模或速度上效率更高，正如我们在开始时承诺的那样。

该过程从大型、预先训练的 LLM 开始，作为教师模型，例如 GPT2 或 LLama。该模型通常非常准确，但需要大量计算资源进行推理。

训练一个更小、更高效的模型（“学生模型”）来模仿教师模型的行为，例如 miniGPT2 或 TinyLlama（尽管 Tinyllama 的构建方式不同）。学生模型从原始训练数据和教师模型生成的输出（软标签）中学习。

以下是从老师GPT2开始的Python师生互动示例：


import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import torch.nn.functional as F# Load the teacher (large) and student (smaller) models
teacher_model_name = "gpt2"  # You can replace this with any large LLM
student_model_name = "tiny-gpt2"  # A smaller variant to act as the student
# Load the teacher model and tokenizer
teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model_name).to("cuda")
teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
# Load the student model and tokenizer
student_model = AutoModelForCausalLM.from_pretrained(student_model_name).to("cuda")
student_tokenizer = AutoTokenizer.from_pretrained(student_model_name)
# Load a dataset for training (e.g., Wikitext for language modeling)
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
# Set training parameters
learning_rate = 5e-5
epochs = 3
optimizer = torch.optim.AdamW(student_model.parameters(), lr=learning_rate)
# Set temperature for softening probabilities
temperature = 2.0
alpha = 0.5  # Weighting factor for combining loss functions
# Training loop for knowledge distillation
for epoch in range(epochs):for i, example in enumerate(dataset):# Get the input textinput_text = example["text"]# Skip empty linesif not input_text.strip():continue# Tokenize the input text for the teacher and student modelsteacher_inputs = teacher_tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=32).to("cuda")student_inputs = student_tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=32).to("cuda")# Get teacher predictions (soft labels)with torch.no_grad():teacher_outputs = teacher_model(**teacher_inputs)teacher_logits = teacher_outputs.logits / temperatureteacher_probs = F.softmax(teacher_logits, dim=-1)# Get student predictionsstudent_outputs = student_model(**student_inputs)student_logits = student_outputs.logits# Calculate distillation loss (Kullback-Leibler divergence)distillation_loss = F.kl_div(input=F.log_softmax(student_logits / temperature, dim=-1),target=teacher_probs,reduction="batchmean",log_target=False) * (temperature ** 2)# Calculate student task loss (Cross-Entropy with true labels)target_labels = student_inputs["input_ids"]task_loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), target_labels.view(-1), ignore_index=student_tokenizer.pad_token_id)# Combined lossloss = alpha * distillation_loss + (1 - alpha) * task_loss# Backpropagation and optimizationoptimizer.zero_grad()loss.backward()optimizer.step()# Print training progressif i % 100 == 0:print(f"Epoch [{epoch + 1}/{epochs}], Step [{i}], Loss: {loss.item():.4f}")
print("Knowledge distillation completed!")

4. 权重共享

通过在多个模型组件之间共享参数，我们可以减少神经网络的内存占用。当部分或所有层共享同一组权重而不是每个层或组件都有唯一的权重时，模型必须保留的参数数量会大大减少。可以先验地定义模型的架构，事先使用共享权重，或者在训练后将权重共享作为模型压缩技术。例如，一种可能性是将权重聚类，如下面的代码所示：


import torch
import numpy as np
from sklearn.cluster import KMeansdef apply_weight_sharing(model, num_clusters=16):# Iterate through each parameter in the modelfor name, param in model.named_parameters():if param.requires_grad:  # Only consider trainable parameters# Flatten the weights into a 1D array for clusteringweights = param.data.cpu().numpy().flatten().reshape(-1, 1)# Apply k-means clusteringkmeans = KMeans(n_clusters=num_clusters)kmeans.fit(weights)# Replace weights with their corresponding cluster centroidscluster_centroids = kmeans.cluster_centers_labels = kmeans.labels_# Map the original weights to their shared valuesshared_weights = np.array([cluster_centroids[label] for label in labels]).reshape(param.data.shape)# Update the model's parameters with the shared weightsparam.data = torch.tensor(shared_weights, dtype=param.data.dtype).to(param.device)return model
# Example usage with a pre-trained model
from transformers import GPT2LMHeadModel
model = GPT2LMHeadModel.from_pretrained("gpt2")
model = apply_weight_sharing(model, num_clusters=16)  # Apply weight sharing with 16 clusters
print("Weight sharing applied to the model!")