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）或者注意力对齐的方法。

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

8. 总结

模型蒸馏是个很有用的技术，特别是现在大模型越来越大，部署成本越来越高。通过蒸馏，我们可以让一个小模型获得接近大模型的能力，同时大大降低部署门槛。

从我这次实践来看，GLM-4-9B蒸馏到3B模型，效果比预想的要好。小模型在大多数任务上能达到教师模型70%-80%的水平，但推理速度提升了3-4倍，显存占用只有原来的三分之一左右。对于很多实际应用场景来说，这个性价比是很高的。

当然，蒸馏也不是万能的。有些特别复杂的任务，小模型确实学不会。而且蒸馏过程比较耗时，需要大量的计算资源和高质量的数据。

如果你也想尝试模型蒸馏，我的建议是从小规模开始。先选一个简单的任务，用少量的数据试试水，看看效果如何。等掌握了基本流程，再逐步扩大规模。蒸馏过程中要多观察、多调整，特别是损失函数的设计和训练参数的设置，对最终效果影响很大。

最后提醒一点，蒸馏后的模型虽然小，但还是要遵守相关的使用规定。特别是如果教师模型有使用限制，学生模型也要遵守同样的限制。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。