大模型知识蒸馏实战：HuggingFace Teacher-Student模型压缩与ONNX导出部署

大模型知识蒸馏是AI模型部署环节的关键技术。通过Teacher-Student架构，将大模型的推理能力迁移到小模型中，在保持90%以上精度的同时将参数量压缩至1/10，显著降低推理延迟和显存占用。本文围绕HuggingFace Transformers的蒸馏流程、蒸馏损失函数设计、ONNX导出量化，给出完整代码实现。

知识蒸馏原理与Teacher-Student架构设计

知识蒸馏的核心思想是让学生模型学习教师模型的软标签（soft targets）输出分布，而非硬标签。软标签包含类间关系信息，提供比one-hot编码更丰富的监督信号。温度参数T控制输出分布的平滑程度，T越高分布越平滑，知识传递越充分。

蒸馏损失由两部分组成：KL散度损失（学生匹配教师的软标签分布）和交叉熵损失（学生匹配真实硬标签）。总损失为两者的加权和：

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class DistillationLoss(nn.Module):
    def __init__(self, alpha=0.5, temperature=4.0):
        super().__init__()
        self.alpha = alpha
        self.temperature = temperature
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_logits, teacher_logits, labels):
        # 软标签蒸馏损失
        soft_student = F.log_softmax(student_logits / self.temperature, dim=-1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=-1)
        distill_loss = self.kl_loss(soft_student, soft_teacher) * (self.temperature ** 2)

        # 硬标签交叉熵损失
        hard_loss = self.ce_loss(student_logits, labels)

        return self.alpha * distill_loss + (1 - self.alpha) * hard_loss

# 加载教师模型（大模型，冻结参数）
teacher_model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-chinese", num_labels=10
)
teacher_model.eval()
for param in teacher_model.parameters():
    param.requires_grad = False

# 加载学生模型（小模型，蒸馏目标）
student_model = AutoModelForSequenceClassification.from_pretrained(
    "hfl/rbt3", num_labels=10  # RoBERTa-3层，比BERT-base少9层
)

教师模型选择BERT-base-chinese（12层，110M参数），学生模型选择rbt3（3层，30M参数）。温度参数T设为4.0，alpha设为0.5，平衡蒸馏损失和真实标签损失。

HuggingFace Transformers蒸馏训练流程实现

使用HuggingFace Trainer实现蒸馏训练，自定义Trainer子类注入教师模型推理：

from transformers import Trainer, TrainingArguments
from torch.utils.data import DataLoader

class DistillationTrainer(Trainer):
    def __init__(self, teacher_model, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model
        self.teacher_model.to(self.model.device)
        self.distill_loss = DistillationLoss(alpha=0.5, temperature=4.0)

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")

        # 教师模型前向推理（不计算梯度）
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
            teacher_logits = teacher_outputs.logits

        # 学生模型前向推理
        student_outputs = model(**inputs)
        student_logits = student_outputs.logits

        loss = self.distill_loss(student_logits, teacher_logits, labels)

        return (loss, student_outputs) if return_outputs else loss

training_args = TrainingArguments(
    output_dir="./distilled-model",
    num_train_epochs=10,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=128,
    learning_rate=5e-5,
    warmup_ratio=0.1,
    weight_decay=0.01,
    logging_steps=100,
    eval_strategy="steps",
    eval_steps=500,
    save_strategy="steps",
    save_steps=1000,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    fp16=True,
)

trainer = DistillationTrainer(
    teacher_model=teacher_model,
    model=student_model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)

trainer.train()
trainer.save_model("./distilled-model/final")

训练中教师模型全程冻结参数，仅前向推理产生软标签。学生模型在软标签和硬标签的双重监督下更新参数。fp16混合精度训练可将显存占用减半，batch_size提升到64。

蒸馏效果评估与中间层特征对齐

仅蒸馏输出层 logits 的效果有限。引入中间层特征对齐（Intermediate Layer Matching），让学生模型在隐层维度上也逼近教师模型。使用MSE损失对齐两者的隐层输出：

class FeatureDistillationTrainer(Trainer):
    def __init__(self, teacher_model, teacher_layer_idx, student_layer_idx, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model
        self.teacher_model.to(self.model.device)
        self.distill_loss = DistillationLoss(alpha=0.5, temperature=4.0)
        self.teacher_layer_idx = teacher_layer_idx
        self.student_layer_idx = student_layer_idx
        # 隐层维度对齐投影层（学生768维 -> 教师768维）
        self.projector = nn.Linear(768, 768).to(self.model.device)

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")

        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs, output_hidden_states=True)
            teacher_hidden = teacher_outputs.hidden_states[self.teacher_layer_idx]

        student_outputs = model(**inputs, output_hidden_states=True)
        student_hidden = student_outputs.hidden_states[self.student_layer_idx]

        # 隐层特征对齐损失
        projected = self.projector(student_hidden)
        feature_loss = F.mse_loss(projected, teacher_hidden)

        # 输出层蒸馏损失
        distill_loss = self.distill_loss(
            student_outputs.logits, teacher_outputs.logits, labels
        )

        total_loss = distill_loss + 0.3 * feature_loss

        return (total_loss, student_outputs) if return_outputs else total_loss

评估指标对比（文本分类任务，10分类）：

特征蒸馏额外提升了0.8个百分点的准确率，推理延迟几乎无变化——投影层仅在训练时使用，推理时丢弃。

ONNX导出与量化部署

蒸馏后的模型导出为ONNX格式，配合ONNX Runtime实现跨平台推理加速。INT8量化进一步压缩模型体积和推理延迟：

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 导出ONNX
dummy_input = tokenizer("测试样本", return_tensors="pt", padding="max_length",
                         max_length=128, truncation=True)
dummy_input = {k: v.to(student_model.device) for k, v in dummy_input.items()}

torch.onnx.export(
    student_model,
    (dummy_input,),
    "./distilled-model/model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size"},
        "attention_mask": {0: "batch_size"},
        "logits": {0: "batch_size"},
    },
    opset_version=17,
)

# 验证ONNX模型
onnx_model = onnx.load("./distilled-model/model.onnx")
onnx.checker.check_model(onnx_model)

# INT8动态量化
quantize_dynamic(
    model_input="./distilled-model/model.onnx",
    model_output="./distilled-model/model_int8.onnx",
    weight_type=QuantType.QInt8,
)

# ONNX Runtime推理
import onnxruntime as ort

sess = ort.InferenceSession("./distilled-model/model_int8.onnx",
                            providers=["CUDAExecutionProvider"])

import numpy as np
input_ids = np.array([[101, 6821, 3221, 102]], dtype=np.int64)
attention_mask = np.array([[1, 1, 1, 1]], dtype=np.int64)

outputs = sess.run(None, {
    "input_ids": input_ids,
    "attention_mask": attention_mask,
})
logits = outputs[0]
predicted = np.argmax(logits, axis=-1)

ONNX INT8量化后的部署效果：

从教师模型110M参数到INT8量化部署，模型体积从440MB压缩到31MB（14:1），推理延迟从85ms降到8ms（10.6:1），精度损失控制在0.9个百分点以内。知识蒸馏的价值在于：用教师模型的软标签知识弥补学生模型容量不足的缺陷，再通过ONNX量化进一步压榨推理效率，实现精度与性能的最佳平衡。