初探大模型压缩

文摘   2024-10-28 07:02   上海  
点击下方卡片,关注“AI生成未来

一般地,语言模型越大越好,改进LLM的方式非常简单: 更多的数据 + 更多的参数 + 更多的计算 = 更好的性能。但是,使用100B + 参数模型存在着明显的挑战。例如,使用 FP16的100B 参数模型仅存储空间就需要200GB!大多数消费设备(如手机、平板电脑、笔记本电脑)无法处理这么大模型。如何能把它们变小呢?

1. 模型压缩

模型压缩的目的是在不牺牲性能的情况下减少机器学习模型的大小。这适用于大型神经网络,因为它们常常过度参数化(即由冗余的计算单元组成)。

模型压缩的主要好处是降低推理成本,这意味着大模型(即在本地笔记本电脑上运行 LLM)更广泛使用,人工智能与消费产品的低成本集成,以及支持用户隐私和安全的设备上推理。

模型压缩技术的范围很广,主要有3大类:

  1. 量化ーー用较低精度的数据类型表示模型

  2. 修剪ーー从模型中删除不必要的组件

  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, f1

5.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


致谢

如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。同时,欢迎给个星标⭐,以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进,共同探索未知,见证一个充满希望和伟大的未来!



技术交流

加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!

AI生成未来
领先的AIGC和具身智能、大模型技术交流社区,关注LLM、CV、深度学习、生成式等AI领域前沿技术
 最新文章