GLM-4-9B模型蒸馏实战：小模型性能提升秘籍-程序员充电站

GLM-4-9B模型蒸馏实战：小模型性能提升秘籍

最近在折腾大模型部署的时候，经常遇到一个头疼的问题：模型太大，显存不够用。特别是像GLM-4-9B这样的模型，虽然性能不错，但动辄需要几十GB的显存，普通显卡根本跑不动。

这时候就有人问了：有没有办法让模型变小一点，但性能别掉太多？还真有，这就是今天要聊的模型蒸馏技术。

简单来说，模型蒸馏就像老师教学生。一个大模型（老师）把自己的知识教给一个小模型（学生），让学生既能保持较小的体积，又能学到老师的大部分能力。我最近用这个方法把GLM-4-9B蒸馏成了一个3B的小模型，效果还挺不错，今天就把整个过程分享给大家。

1. 准备工作：理解蒸馏的核心思路

在开始动手之前，咱们先搞清楚蒸馏到底在做什么。传统的训练是让模型直接学习任务，比如分类、生成文本。而蒸馏是让一个小模型去模仿一个大模型的输出。

这里有个关键点：蒸馏不只是模仿最终的结果，更重要的是模仿大模型在生成每个词时的“思考过程”。大模型在生成文本时，会对每个可能的词有一个概率分布，这个分布包含了丰富的知识。比如在生成“今天天气”后面接什么词时，大模型可能觉得“很好”概率0.4，“不错”概率0.3，“晴朗”概率0.2，而小模型如果只学最终输出的“很好”，就丢失了很多信息。

蒸馏就是要让小模型学会这个完整的概率分布，这样它就能理解为什么“很好”比“晴朗”更合适，而不仅仅是记住答案。

2. 数据采集：准备“教材”

蒸馏的第一步是准备训练数据。这里的数据不是普通的标注数据，而是大模型在各种输入上的输出。你可以把它想象成老师备课的教案。

我用的方法是准备一些高质量的对话数据，然后让GLM-4-9B来生成回答。这里有个小技巧：不要只用一种提问方式，要多样化。比如同一个问题，可以用不同的表述方式，或者从不同的角度提问，这样能让小模型学到更全面的知识。

import json from transformers import AutoTokenizer, AutoModelForCausalLM import torch # 加载GLM-4-9B作为教师模型 teacher_model_name = "THUDM/glm-4-9b-chat" tokenizer = AutoTokenizer.from_pretrained(teacher_model_name, trust_remote_code=True) teacher_model = AutoModelForCausalLM.from_pretrained( teacher_model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True ) # 准备一些示例问题 questions = [ "请解释一下什么是机器学习", "写一个简单的Python函数计算斐波那契数列", "如何快速学习一门新的编程语言", "用三个关键词总结人工智能的发展趋势", "给初学者推荐学习深度学习的路线" ] # 收集教师模型的输出 training_data = [] for question in questions: # 构建对话格式 messages = [{"role": "user", "content": question}] # 使用apply_chat_template处理对话 inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True ) # 生成回答 with torch.no_grad(): outputs = teacher_model.generate( **inputs.to(teacher_model.device), max_new_tokens=500, do_sample=True, temperature=0.7, top_p=0.9 ) # 解码输出 response = tokenizer.decode(outputs[0], skip_special_tokens=True) # 提取生成的回答部分 # 这里需要根据具体的对话模板来调整 full_text = response # 假设我们只保留模型生成的部分 generated_text = full_text.split(question)[-1].strip() training_data.append({ "question": question, "teacher_response": generated_text, "full_output": response }) # 保存数据 with open("distillation_data.json", "w", encoding="utf-8") as f: json.dump(training_data, f, ensure_ascii=False, indent=2)

这段代码会生成一个包含教师模型回答的数据集。在实际操作中，你可能需要准备几千甚至几万条这样的数据，覆盖不同的领域和问题类型。

3. 学生模型选择：找个好“苗子”

选学生模型是个技术活。不能随便找个小模型就开始蒸馏，得考虑几个因素：

首先，学生模型的架构最好和教师模型相似。GLM-4-9B用的是GLM架构，所以我也选了一个GLM架构的小模型作为起点。如果架构差异太大，蒸馏效果可能会打折扣。

其次，要考虑学生模型的容量。太小了学不会，太大了又失去了蒸馏的意义。我选了一个3B参数的GLM模型，这个大小在消费级显卡上就能跑，而且有足够的容量来学习教师模型的知识。

最后，还要看学生模型的基础能力。如果学生模型本身太差，就像让小学生直接学大学课程，效果肯定不会好。我选了一个在通用任务上表现还不错的3B模型作为基础。

from transformers import AutoModelForCausalLM # 加载学生模型（这里以一个小型GLM模型为例） student_model_name = "你的小型GLM模型路径" student_model = AutoModelForCausalLM.from_pretrained( student_model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True ) # 检查模型参数 print(f"学生模型参数量: {student_model.num_parameters() / 1e9:.1f}B") print(f"教师模型参数量: {teacher_model.num_parameters() / 1e9:.1f}B")

4. 损失函数设计：怎么“教”最有效

蒸馏的核心在于损失函数的设计。传统的训练只用真实标签，而蒸馏要用到教师模型的软标签（概率分布）。

我用了两种损失函数的组合：

第一种是蒸馏损失，让学生模型的输出概率分布尽量接近教师模型。这里用KL散度来衡量两个分布的差异。

第二种是任务损失，让学生模型也能直接学习任务。虽然蒸馏是主要目标，但也不能完全忽略原始任务。

import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): def __init__(self, alpha=0.7, temperature=2.0): super().__init__() self.alpha = alpha # 蒸馏损失的权重 self.temperature = temperature # 温度参数 self.ce_loss = nn.CrossEntropyLoss() def forward(self, student_logits, teacher_logits, labels=None): """ student_logits: 学生模型的输出logits [batch, seq_len, vocab_size] teacher_logits: 教师模型的输出logits [batch, seq_len, vocab_size] labels: 真实标签（可选） """ batch_size, seq_len, vocab_size = student_logits.shape # 计算蒸馏损失（KL散度） # 使用温度缩放 student_probs = F.log_softmax(student_logits / self.temperature, dim=-1) teacher_probs = F.softmax(teacher_logits / self.temperature, dim=-1) # KL散度损失 kd_loss = F.kl_div( student_probs.view(-1, vocab_size), teacher_probs.view(-1, vocab_size), reduction='batchmean' ) * (self.temperature ** 2) total_loss = kd_loss # 如果有真实标签，加入任务损失 if labels is not None: # 调整labels形状以匹配logits if labels.dim() == 1: labels = labels.unsqueeze(0) # 计算交叉熵损失 ce_loss = self.ce_loss( student_logits.view(-1, vocab_size), labels.view(-1) ) # 加权组合 total_loss = self.alpha * kd_loss + (1 - self.alpha) * ce_loss return total_loss

温度参数在这里很关键。温度高的时候，概率分布更平滑，小模型能学到更多细节；温度低的时候，分布更尖锐，更关注主要的预测。我一般从温度2.0开始，然后根据效果调整。

5. 训练流程：手把手教你蒸馏

有了数据、模型和损失函数，就可以开始训练了。蒸馏训练和普通训练有些不同，需要同时用到教师模型和学生模型。

from torch.utils.data import Dataset, DataLoader from tqdm import tqdm import torch.optim as optim class DistillationDataset(Dataset): def __init__(self, data_file, tokenizer, max_length=512): with open(data_file, 'r', encoding='utf-8') as f: self.data = json.load(f) self.tokenizer = tokenizer self.max_length = max_length def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] # 编码输入 input_text = item["question"] teacher_output = item["teacher_response"] # 构建完整的输入输出 full_text = input_text + "\n" + teacher_output # 分词 encoding = self.tokenizer( full_text, truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt" ) return { "input_ids": encoding["input_ids"].squeeze(), "attention_mask": encoding["attention_mask"].squeeze(), "teacher_text": teacher_output } def train_distillation(): # 准备数据 dataset = DistillationDataset("distillation_data.json", tokenizer) dataloader = DataLoader(dataset, batch_size=2, shuffle=True) # 初始化损失函数 criterion = DistillationLoss(alpha=0.7, temperature=2.0) # 优化器 optimizer = optim.AdamW(student_model.parameters(), lr=5e-5) # 训练循环 num_epochs = 3 student_model.train() teacher_model.eval() # 教师模型不更新参数 for epoch in range(num_epochs): total_loss = 0 progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}") for batch in progress_bar: input_ids = batch["input_ids"].to(student_model.device) attention_mask = batch["attention_mask"].to(student_model.device) # 前向传播（学生模型） student_outputs = student_model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=False, use_cache=False ) student_logits = student_outputs.logits # 前向传播（教师模型） with torch.no_grad(): teacher_outputs = teacher_model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=False, use_cache=False ) teacher_logits = teacher_outputs.logits # 计算损失 # 注意：这里我们使用输入作为标签（自回归语言模型） # 实际应用中可能需要调整 shift_logits = student_logits[..., :-1, :].contiguous() shift_labels = input_ids[..., 1:].contiguous() shift_teacher_logits = teacher_logits[..., :-1, :].contiguous() loss = criterion( shift_logits.view(-1, shift_logits.size(-1)), shift_teacher_logits.view(-1, shift_teacher_logits.size(-1)), shift_labels.view(-1) ) # 反向传播 optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(student_model.parameters(), 1.0) optimizer.step() total_loss += loss.item() progress_bar.set_postfix({"loss": loss.item()}) avg_loss = total_loss / len(dataloader) print(f"Epoch {epoch+1} 平均损失: {avg_loss:.4f}") # 保存蒸馏后的模型 student_model.save_pretrained("distilled_glm_3b") tokenizer.save_pretrained("distilled_glm_3b")

训练过程中有几个注意事项：

第一，批次大小不能太大。因为要同时跑教师模型和学生模型，显存占用差不多是两者的和。我用的批次大小是2，如果你的显卡更好，可以适当调大。

第二，学习率要调低一点。蒸馏训练相对精细，学习率太大会破坏学生模型已经学到的知识。

第三，训练轮数不用太多。一般3-5个epoch就够了，过度训练反而可能导致学生模型过度拟合教师模型的特定输出。

6. 效果评估：看看“学生”学得怎么样

训练完成后，最重要的一步是评估效果。不能只看损失函数的值，要看实际生成的效果。

我用了几个方法来评估：

首先是生成质量对比。用同样的输入让教师模型和学生模型都生成回答，然后对比两者的差异。这里不仅要看内容是否相关，还要看流畅度、逻辑性等。

def compare_generation(question, max_length=200): """对比教师模型和学生模型的生成效果""" # 准备输入 messages = [{"role": "user", "content": question}] # 教师模型生成 print("=== 教师模型生成 ===") teacher_inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True ) with torch.no_grad(): teacher_outputs = teacher_model.generate( **teacher_inputs.to(teacher_model.device), max_new_tokens=max_length, do_sample=True, temperature=0.7 ) teacher_response = tokenizer.decode(teacher_outputs[0], skip_special_tokens=True) print(teacher_response) print() # 学生模型生成 print("=== 学生模型生成 ===") student_inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True ) with torch.no_grad(): student_outputs = student_model.generate( **student_inputs.to(student_model.device), max_new_tokens=max_length, do_sample=True, temperature=0.7 ) student_response = tokenizer.decode(student_outputs[0], skip_special_tokens=True) print(student_response) # 测试几个问题 test_questions = [ "什么是深度学习？", "用Python写一个快速排序算法", "解释一下注意力机制的原理" ] for q in test_questions: print(f"\n问题: {q}") compare_generation(q) print("="*50)

其次是性能测试。在同样的硬件上，测试两个模型的推理速度、显存占用等。这是蒸馏的主要目的之一，要让小模型在资源有限的情况下也能用。

import time import psutil import GPUtil def benchmark_model(model, tokenizer, prompt, num_runs=10): """基准测试模型性能""" # 准备输入 messages = [{"role": "user", "content": prompt}] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True ) inputs = inputs.to(model.device) # 预热 with torch.no_grad(): _ = model.generate(**inputs, max_new_tokens=10) # 测试生成速度 start_time = time.time() for _ in range(num_runs): with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=100, do_sample=False # 为了测试速度，关闭采样 ) end_time = time.time() avg_time = (end_time - start_time) / num_runs # 检查显存使用 gpus = GPUtil.getGPUs() gpu_memory = gpus[0].memoryUsed if gpus else 0 return { "avg_generation_time": avg_time, "gpu_memory_mb": gpu_memory, "output_length": len(outputs[0]) } # 测试性能 test_prompt = "请简要介绍人工智能的发展历史" print("教师模型性能:") teacher_stats = benchmark_model(teacher_model, tokenizer, test_prompt) print(f"平均生成时间: {teacher_stats['avg_generation_time']:.2f}秒") print(f"GPU显存使用: {teacher_stats['gpu_memory_mb']} MB") print("\n学生模型性能:") student_stats = benchmark_model(student_model, tokenizer, test_prompt) print(f"平均生成时间: {student_stats['avg_generation_time']:.2f}秒") print(f"GPU显存使用: {student_stats['gpu_memory_mb']} MB") # 计算提升比例 speedup = teacher_stats['avg_generation_time'] / student_stats['avg_generation_time'] memory_saving = teacher_stats['gpu_memory_mb'] / student_stats['gpu_memory_mb'] print(f"\n性能对比:") print(f"速度提升: {speedup:.1f}倍") print(f"显存节省: {memory_saving:.1f}倍")

最后是任务评估。在一些标准任务上测试两个模型的表现，比如常识推理、代码生成、文本摘要等。虽然小模型不可能在所有任务上都达到大模型的水平，但在大多数任务上应该能有不错的表现。

7. 实战技巧与常见问题

在实际操作中，我遇到了一些问题，也总结了一些经验，分享给大家：

数据质量是关键。如果教师模型生成的数据质量不高，学生模型学到的就是垃圾。建议先用一些筛选机制，比如只保留置信度高的输出，或者人工检查一部分数据。

温度参数需要调优。温度太高，学生模型学不到重点；温度太低，学生模型可能过度拟合。我一般会尝试几个不同的温度值（1.5, 2.0, 3.0），然后选效果最好的。

注意灾难性遗忘。学生模型在学教师模型知识的同时，可能会忘记自己原本的能力。可以在损失函数中加入对学生模型原始输出的约束，或者用多任务学习的方式。

逐步蒸馏。如果直接从一个很大的模型蒸馏到很小的模型效果不好，可以尝试分步蒸馏：大模型→中模型→小模型。每一步的差距不要太大。

利用中间层特征。除了最终的输出分布，教师模型的中间层特征也包含很多信息。可以让学生模型的某些层去匹配教师模型对应层的特征，这叫做特征蒸馏。

class FeatureDistillationLoss(nn.Module): """特征蒸馏损失，匹配中间层特征""" def __init__(self, layer_mapping): """ layer_mapping: 字典，指定学生模型的哪些层对应教师模型的哪些层 例如：{0: 2, 2: 5, 4: 8} 表示学生第0层对应教师第2层 """ super().__init__() self.layer_mapping = layer_mapping self.mse_loss = nn.MSELoss() def forward(self, student_features, teacher_features): """ student_features: 学生模型各层的特征列表 teacher_features: 教师模型各层的特征列表 """ loss = 0 for s_layer, t_layer in self.layer_mapping.items(): if s_layer < len(student_features) and t_layer < len(teacher_features): # 对齐特征维度（如果需要） s_feat = student_features[s_layer] t_feat = teacher_features[t_layer] # 如果维度不匹配，可以用线性层投影 if s_feat.shape[-1] != t_feat.shape[-1]: # 这里简化处理，实际可能需要更复杂的对齐 continue loss += self.mse_loss(s_feat, t_feat) return loss

处理输出长度差异。教师模型和学生模型的输出长度可能不同，需要对齐。可以用动态时间规整（DTW）或者注意力对齐的方法。

注意计算效率。蒸馏训练的计算成本很高，因为要同时跑两个模型。可以考虑用梯度累积来增大有效批次大小，或者用混合精度训练来节省显存。