初探大模型压缩
初探大模型压缩
【引】感谢大家对联想AIPC的关注！大模型在智能终端上应用使AI更方便地服务于我们的工作和生活，《从苹果智能看端上大模型应用》为我们提供了参考。
一般地，语言模型越大越好，改进LLM的方式非常简单: 更多的数据 + 更多的参数 + 更多的计算 = 更好的性能。但是，使用100B + 参数模型存在着明显的挑战。例如，使用 FP16的100B 参数模型仅存储空间就需要200GB！大多数消费设备(如手机、平板电脑、笔记本电脑)无法处理这么大模型。如何能把它们变小呢？
1. 模型压缩
模型压缩的目的是在不牺牲性能的情况下减少机器学习模型的大小。这适用于大型神经网络，因为它们常常过度参数化(即由冗余的计算单元组成)。
模型压缩的主要好处是降低推理成本，这意味着大模型(即在本地笔记本电脑上运行 LLM)的更广泛使用，人工智能与消费产品的低成本集成，以及支持用户隐私和安全的设备上推理。
模型压缩技术的范围很广，主要有3大类：
- 量化ーー用较低精度的数据类型表示模型
- 修剪ーー从模型中删除不必要的组件
- 知识蒸馏ーー用大模型训练小模型
这些方法是相互独立的。因此，来自多个类别的技术组合在一起可以获得最大的压缩。
2. 量化
虽然量化听起来像一个可怕而复杂的词，但它是一个简单的想法，主要是降低模型参数的精度。我们可以把这看作是在保持图片核心属性的同时，将高分辨率图像转换为低分辨率图像。
两种常见的量化技术是训练后量化(PTQ)和量化感知训练(QAT)。
2.1 训练后量化(PTQ)
给定一个神经网络，后训练量化(PTQ)通过用低精度数据类型(例如 FP16到 INT-8)替换参数来压缩模型。这是减少模型计算需求的最快和最简单的方法之一，因为它不需要额外的训练或数据标注。
虽然这是一种相对容易的削减模型成本的方法，但这种方法中过多的量化(例如，FP16到 INT4)常常会导致性能下降，从而限制了 PTQ 的潜在收益。
2.2量化感知训练
对于需要更大压缩的情况，PTQ 的局限性可以通过使用低精度数据类型的训练模型(从头开始)来克服。这就是量化感知训练(QAT)的背后思想。虽然这种方法在技术上要求更高，但它可以产生一个更小、性能更好的模型。例如，BitNet 体系结构使用三元数据类型(即1.58位)来匹配原始 Llama LLM 的性能。
当然，PTQ 和从头开始的 QAT 之间存在很大的技术差距。两者之间的一种方法是量化感知微调，它包括量化后预训练模型的额外训练。
3. 修剪
修剪的目的是删除对性能影响很小的模型组件，其有效性在于机器学习模型(尤其是大模型)倾向于学习冗余和嘈杂的结构。这里的比喻就像是从树上剪下枯枝，剪枝可以在不伤害树的情况下减小树的体积。修剪方法可以分为两类: 非结构化修剪和结构化修剪。
3.1 非结构化修剪
非结构化剪枝从神经网络中移除不重要的权重(即将它们设置为零)。例如，通过估计对损失函数的影响来计算网络中每个参数的显著性得分。去除具有最小绝对值的权重的方法，由于其简单性和可伸缩性而变得流行起来。
虽然非结构化剪枝的粒度可以显著减少参数计数，但是这些增益一般需要专门的硬件来实现。非结构化剪枝导致稀疏矩阵运算 ，标准硬件往往无法更有效地完成。
3.2 结构化修剪
结构化修剪从神经网络中移除整个结构(例如注意力头，神经元和层)。这避免了专用矩阵运算的问题，因为整个矩阵可以从模型中删除，而不是单独的参数。虽然有各种方法可以识别结构进行修剪，但原则上，它们都试图删除对性能影响最小的结构。
4.知识蒸馏
知识蒸馏将知识从(较大的)教师模型转移到(较小的)学生模型。做到这一点的一种方法是用教师模型生成预测，并用它们来训练学生模型。学习教师模型的输出 logits (即，所有可能的下一个令牌的概率)提供了比原始训练数据更丰富的信息，这提高了学生模型的性能。
最近的蒸馏应用程序完全放弃了对 logit 的需要，而是从教师模型生成的合成数据中学习。一个流行的例子是斯坦福大学的 Alpaca 模型，该模型使用 OpenAI 的 text-davinci-003(即原始 ChatGPT 模型)的合成数据对 LLaMa 7B 模型进行了微调，使其能够遵循用户指令。
5. 实验：用知识蒸馏 + 量化压缩文本分类器
作为一个实验，我们将压缩一个100M 参数模型，该模型将 URL 分类为安全还是不安全(即是否是钓鱼网站)。首先利用知识精馏将100M 参数模型压缩为50M 参数模型。然后，使用4位量化，进一步减少了3倍的内存占用，导致最终的模型是原始模型的1/8。
5.1 环境构建
我们首先导入一些需要使用的库。
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("llmc/phishing-site-classification")
5.2 加载教师模型
加载教师模型，为了帮助加快训练速度，我们需要使用GPU处理器。
# use Nvidia GPU
device = torch.device('cuda')
# Load teacher model and tokenizer
model_path = "llmc/bert-phishing-classifier_teacher"
tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)
.to(device)这个教师模型是 Goolge 的 bert-base-uncase 的一个微调版本，它对钓鱼网站的 URL 执行二进制分类。
5.3 构建学生模型
对于学生模型，需要从头开始初始化，通过从剩余的层中移除两个层和四个注意头来修改模型的架构。
# Load student model
my_config = DistilBertConfig(n_heads=8, n_layers=4) # drop 4 heads per layer and 2 layers
student_model = DistilBertForSequenceClassification
.from_pretrained("distilbert-base-uncased",
config=my_config,)
.to(device)在训练学生模型之前，我们需要对数据集进行标记。这一点很重要，因为模型期望以特定的方式表示输入文本。
在这里，基于每批最长的示例填充，允许将批次表示为 PyTorch 张量。
# define text preprocessing
def preprocess_function(examples):
return tokenizer(examples["text"], padding='max_length', truncation=True)
# tokenize all datasetse
tokenized_data = data.map(preprocess_function, batched=True)
tokenized_data.set_format(type='torch',
columns=['input_ids', 'attention_mask', 'labels'])训练前的另一个重要步骤是在训练期间为模型定义一个评估策略。下面，定义一个函数，它计算给定模型和数据集的准确率、精确率、召回率和 F1得分。
# Function to evaluate model performance
def evaluate_model(model, dataloader, device):
model.eval() # Set model to evaluation mode
all_preds = []
all_labels = []
# Disable gradient calculations
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)
# Forward pass to get logits
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Get predictions
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
# Calculate evaluation metrics
accuracy = accuracy_score(all_labels, all_preds)
precision, recall, f1, _ = precision_recall_fscore_support(all_labels,
all_preds,
average='binary')
return accuracy, precision, recall, f15.4 训练学生模型
为了学生模型同时学习训练集中的可信数据标签(即硬目标)和教师模型的逻辑(即软目标) ，我们必须构造一个特殊的损失函数来考虑两个目标。这是通过将学生和教师的输出概率分布的 KL 散度与学生 logit 的交叉熵损失和基本真理相结合来完成的。
# Function to compute distillation and hard-label loss
def distillation_loss(student_logits, teacher_logits,
true_labels, temperature, alpha):
# Compute soft targets from teacher logits
soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)
# KL Divergence loss for distillation
distill_loss = nn.functional.kl_div(student_soft,
soft_targets,
reduction='batchmean') * (temperature ** 2)
# Cross-entropy loss for hard labels
hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)
# Combine losses
loss = alpha * distill_loss + (1.0 - alpha) * hard_loss
return loss接下来，定义超参数、优化器、训练数据集和测试数据集。
# hyperparameters
batch_size = 32
lr = 1e-4
num_epochs = 5
temperature = 2.0
alpha = 0.5
# define optimizer
optimizer = optim.Adam(student_model.parameters(), lr=lr)
# create training data loader
dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)
# create testing data loader
test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)最后，我们使用 PyTorch 训练学生模型。
# put student model in train mode
student_model.train()
# train model
for epoch in range(num_epochs):
for batch in dataloader:
# Prepare inputs
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
# Disable gradient calculation for teacher model
with torch.no_grad():
teacher_outputs = teacher_model(input_ids,
attention_mask=attention_mask)
teacher_logits = teacher_outputs.logits
# Forward pass through the student model
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} completed with loss: {loss.item()}")
# Evaluate the teacher model
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}")
# Evaluate the student model
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")
# put student model back into train mode
student_model.train()5.5 模型评估
我们可以在独立的验证集上评估模型，也就是说，使用那些不用于训练模型参数或调整超参数的数据。
# create testing data loader
validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)
# Evaluate the teacher model
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}")
# Evaluate the student model
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}")5.6 模型量化
我们再使用 QLoRA 文章中描述的4位 NormalFloat 数据类型和用于计算的 bfloat16设置配置来存储模型参数。
from transformers import BitsAndBytesConfig
# load model in model 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=device,
quantization_config=nf4_config)然后，可以在验证集上评估我们的量化模型。
# Evaluate the student model
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}")压缩之后性能有了小小的提高，一个直观的解释是 Occam 的剃刀原理，该原理指出，简单的模型更好。在这个实验中，模型可能过度参数化了这个二进制分类任务。因此，简化模型可以获得更好的性能。
一句话小结
虽然LLM在各种任务中表现出了令人印象深刻的性能，但是它们在部署到现实世界环境中时存在挑战，模型压缩技术(量化、修剪和知识蒸馏) 通过降低 LLM 计算成本来帮助缓解这些挑战。
【参考资料与关联阅读】
- A Survey of Model Compression and Acceleration for Deep Neural Networks，https://arxiv.org/abs/1710.09282
- A Survey on Model Compression for Large Language Models，https://arxiv.org/abs/2308.07633
- To prune, or not to prune: exploring the efficacy of pruning for model compression，https://arxiv.org/abs/1710.01878
- Distilling the Knowledge in a Neural Network，https://arxiv.org/abs/1503.02531
腾讯云开发者
