利用云端 GPU 动手实践：从零开始蒸馏小模型

2025年4月22日

随着大语言模型（LLM）的不断发展，如何高效地将大模型的能力“压缩”进小模型成为一个热门课题。知识蒸馏（Knowledge Distillation） 正是实现这一目标的重要手段。它允许我们在有限资源（如单卡 GPU）下，让小模型学习大模型的知识，从而达到在推理速度和部署成本上的双赢。

DeepSeek 团队基于 800k 高质量样本对 Qwen/Llama 等小模型进行蒸馏（论文链接），验证了几个关键结论：

蒸馏有效迁移推理能力至小模型，蒸馏模型在多项推理任务中表现出色。
蒸馏比强化学习更高效。在相同资源下，蒸馏方法比直接对小模型进行强化学习训练更能提高性能，其训练成本更低。

本篇博客将带你从零开始，基于开源工具和公有云 GPU 的算力资源，动手实践一个完整的知识蒸馏流程。我们以 DeepSeek-R1 模型生成的回答作为教师知识，训练微软开源的小模型 phi-3-mini-4k-instruct 来模仿它的行为。

本文实践环节是基于 How to distill Deepseek-R1: A Comprehensive Guide 介绍如何在 Spader.AI 上提供的一张 A100 40G 单卡的全训练过程，完整的代码可以从 GitHub 或 Gitee 上下载。

一、准备数据：从指令到对话格式

我们使用由 DeepSeek-R1 生成的推理问答数据集 Magpie-Reasoning-V2-250K-CoT-Deepseek-R1-Llama-70B，其特点是涵盖了 Chain-of-Thought（思维链）风格的详细解题过程，非常适合用于训练小模型学习复杂推理结构。这个数据集包含大量由 Deepseek-R1（基于Llama-70B）模型生成的指令-回答对，以及详细的链式思考过程（Chain-of-Thought）。我们通过 load_dataset 加载其中一部分数据，并对格式进行转换：

from datasets import load_dataset

# 加载由教师模型生成的 Q&A 数据集（取训练集的前20%以加快演示）
dataset = load_dataset(
    "Magpie-Align/Magpie-Reasoning-V2-250K-CoT-Deepseek-R1-Llama-70B",
    split="train[:20%]"
    token="<你的HuggingFace访问令牌>"
)

接着，我们将其统一格式化为对话形式（添加特殊标记）：

def format_instruction(example):
    return {
        "text": (
            "<|user|>\n"
            f"{example['instruction']}\n"
            "<|end|>\n"
            "<|assistant|>\n"
            f"{example['response']}\n"
            "<|end|>"
        )
    }

formatted_dataset = dataset.map(format_instruction, batched=False, remove_columns=dataset.column_names)
formatted_dataset = formatted_dataset.train_test_split(test_size=0.1)  # 划分90%训练、10%验证

二、加载学生模型与分词器

我们选用的是 phi-3-mini-4k-instruct 作为学生模型，它性能轻量、指令对齐能力较好，适合快速蒸馏实践。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "microsoft/phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

# 添加特殊标记用于链式思考的标记符
CUSTOM_TOKENS = ["<think>", "</think>"]
tokenizer.add_special_tokens({"additional_special_tokens": CUSTOM_TOKENS})
tokenizer.pad_token = tokenizer.eos_token

# 加载学生模型（因模型含自定义代码，需 trust_remote_code）
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    device_map="auto",          # 自动将模型加载到 GPU
    torch_dtype=torch.float16,  # 使用 16 位精度减少显存占用
    attn_implementation="flash_attention_2"  # 使用 Flash Attention 2 加速
)
model.resize_token_embeddings(len(tokenizer))  # 模型词表扩充以适配新添加的特殊标记

代码解析：我们首先加载了分词器，并增添特殊标记 <think> 和 </think> 用于表示思考过程的起止。这是因为教师模型的回答中包含了思考步骤，我们希望学生模型也能学习这种链式推理格式。添加特殊标记后，需要调用 resize_token_embeddings 调整模型的词表大小，否则模型无法识别新标记。接着，用 AutoModelForCausalLM.from_pretrained 加载学生模型权重，并使用 device_map="auto" 将模型自动加载到可用 GPU 上（如果租用的是单卡 GPU，则模型会加载到该卡上）。torch_dtype=torch.float16 表示采用半精度浮点数，加速推理和训练。这里还启用了 flash_attention_2 来优化注意力计算（需要硬件支持），进一步提升训练效率。

说明：在知识蒸馏中，我们通常不需要加载教师模型的权重，因为我们使用的是教师模型预先生成的“答案”作为训练目标。这避免了在云端加载一个超大模型（70B 参数或更大）的开销，仅需要较小的学生模型即可完成训练。

三、配置 LoRA：高效微调的关键

由于学生模型（Phi-3）本身仍有 3.8B 参数，直接训练所有参数需要较大的 GPU 显存和计算。我们引入 LoRA (Low-Rank Adaptation) 技术对模型进行高效微调，只训练很小一部分参数，以降低资源占用。LoRA 通过为模型的权重增加低秩矩阵的偏置来进行训练，冻结原始权重不变。以下代码配置了 LoRA 的超参数：

from peft import LoraConfig

peft_config = LoraConfig(
    r=8,                # 低秩矩阵的秩 (rank)
    lora_alpha=16,      # LoRA的缩放系数
    lora_dropout=0.2,   # LoRA层的Dropout概率
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 应用LoRA的目标层
    bias="none",
    task_type="CAUSAL_LM"
)

这里我们将 LoRA 的秩设为 8，表示每个被微调的权重矩阵将增加两个维度为 (原维度, 8) 和 (8, 原维度) 的低秩矩阵；lora_alpha=16 是缩放因子，会在应用到模型时乘到 LoRA 的增量上；lora_dropout=0.2 则在训练中对 LoRA 分支施加20%的dropout以增强泛化。target_modules 指定只对模型中 Q、K、V、O 投影层应用 LoRA（这些是 Transformer 自注意力机制中的关键矩阵，参数量大且对模型输出影响大）。通过这样的配置，模型大部分参数保持冻结不动，仅有少量新引入的参数参与训练，使得在云端 GPU 上以较小显存开销即可完成模型微调。

提示：LoRA 大幅减少了训练所需的参数量。例如 Phi-3 模型原本有 3.8B 参数，引入 LoRA 后实际训练的参数仅占不到 1%，这意味着显存占用和计算量也大幅降低，非常适合在租用的单 GPU 环境中运行长时间训练。

四、数据预处理与整理

使用 DataCollatorForLanguageModeling 自动处理 padding 和标签构建。

from transformers import DataCollatorForLanguageModeling

# 数据整理器：为因果语言建模准备批次数据（不使用掩码语言模型）
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=8  # 将长度填充对齐到8的倍数，提高 GPU 利用率
)

DataCollatorForLanguageModeling 会将我们的文本样本转换成模型可训练的张量批次，并自动进行填充（padding）。设置 mlm=False 表示我们在做自回归语言模型训练：模型需要根据前面的文本预测下一个词。由于我们的格式中用户提问和助手回答都在同一序列里，在训练时模型会尝试预测包括助手回答在内的整段文本序列。这实际上让学生模型学习在看到<|assistant|>标志后，输出正确的回答序列（而 <|user|> 部分作为上下文提供，不作为需要预测的目标）。最终的损失函数就是标准的交叉熵损失：比较学生模型在每个时间步生成下一个词的概率与教师答案实际词的 One-Hot 分布，计算平均负对数似然。

五、训练参数配置

我们设置合理的 batch size、学习率、训练轮数，适配云端单卡GPU训练条件。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./phi-3-deepseek-finetuned-checkpoints",  # 模型输出保存路径
    num_train_epochs=2,             # 训练轮数（Epoch数）
    per_device_train_batch_size=4,  # 每块 GPU 上的训练批大小
    per_device_eval_batch_size=2,   # 每块 GPU 上的验证批大小
    gradient_accumulation_steps=4,  # 梯度累积步数
    evaluation_strategy="epoch",    # 每个 Epoch 结束后评估验证集
    save_strategy="epoch",          # 每个 Epoch 结束后保存模型
    logging_strategy="steps", 
    logging_steps=50,               # 每 50 步打印一次日志
    learning_rate=2e-5,             # 初始学习率
    bf16=True,                      # 使用 bfloat16 混合精度训练 (适用于 Ampere 架构 GPU)
    optim="paged_adamw_32bit",      # 优化器：使用 bitsandbytes 的 32 位 AdamW 优化内存
    max_grad_norm=0.3,              # 最大梯度范数（用于梯度裁剪）
    warmup_ratio=0.03,              # 3% 训练步骤用于学习率预热
    lr_scheduler_type="cosine",     # 余弦退火学习率调度
    report_to="none"   
)

参数解读：上述配置中，per_device_train_batch_size=4 和 gradient_accumulation_steps=4 结合，意味着实际每次参数更新相当于 16 个样本（4×4）的效果，有助于提高稳定性同时照顾显存限制；evaluation_strategy="epoch"让我们每轮结束后计算验证集损失，便于观察学生模型进展；bf16=True利用了 Ampere 架构 GPU 的 BFloat16 精度来替代 FP16，提供更宽的动态范围（训练更稳定）。此外，我们选择了 bitsandbytes 库提供的 paged_adamw_32bit 优化器来高效利用内存，并设置了小的 max_grad_norm 防止梯度爆炸。学习率调度使用余弦退火策略并配合 3% 的热身（warmup）。

六、启动蒸馏训练

有了以上准备，我们使用 TRL库 提供的 SFTTrainer 来封装训练流程。SFTTrainer（Supervised Fine-Tuning Trainer）是对 Hugging Face Trainer 的扩展，简化了微调大型语言模型（包括 LoRA 等 PEFT 方法）的流程：

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    data_collator=data_collator,
    peft_config=peft_config
)

trainer.train()
trainer.save_model(training_args.output_dir)
tokenizer.save_pretrained(training_args.output_dir)

本次训练是取 20% 的训练数据和 5 个 epoch，在 Spader.AI 上单卡 A100 40G 花了 10 个小时，其结果如下，详细的训练日志上传到 GitHub 。

Epoch	Step	Training Loss	Validation Loss
1	2811	0.4385	0.4408
2	5622	0.4266	0.4272
3	8433	0.4224	0.4215
4	11244	0.4216	0.4194
5	14055	0.4291	0.4190

training result

从结果看整个训练过程平稳，损失波动较小，没有剧烈的波动情况出现，这显示训练配置合理，学习率适合。Validation Loss 始终低于或接近 Training Loss，这说明模型泛化表现良好，没有明显过拟合，训练和验证损失差距较小，表现理想。如果时间允许还可以继续训练几个 epoch。

七、合并并保存模型

final_model = trainer.model.merge_and_unload()
final_model.save_pretrained("./phi-3-deepseek-finetuned")
tokenizer.save_pretrained("./phi-3-deepseek-finetuned")

八、推理测试：看看小模型学到了什么？

使用以下代码快速验证模型输出是否符合预期风格：

from transformers import pipeline

pipe = pipeline("text-generation", model=training_args.output_dir, tokenizer=tokenizer, device_map="auto")
query = "<|user|>\nWhat's the probability of rolling a 7 with two dice?\n<|end|>\n<|assistant|>\n"
output = pipe(query, max_new_tokens=5000, do_sample=True, temperature=0.7, eos_token_id=tokenizer.eos_token_id)
print(output[0]["generated_text"])

✅ 小模型的回答通常会模仿大模型即此案例中的 DeepSeek 的链式推理结构，即在<think></think>之间展示思考过程列举骰子组合并说明为什么答案是 1/6，展现出明确的逻辑过程。如截图所示

总结

在本教程中，我们通过一个实战案例学习了：

如何使用公有云 GPU 执行完整的知识蒸馏流程；
如何加载并构建一个蒸馏任务的数据管道；
如何用 LoRA 技术高效微调小模型；
如何训练并评估学生模型效果；

这种从“生成数据 → 蒸馏训练 → 模型部署”的流程，是未来 LLM 小模型落地的重要方向。我们可以动手实践更进一步理解了知识蒸馏的意义，并能将其应用到自己的项目中。

一、准备数据：从指令到对话格式​

二、加载学生模型与分词器​

三、配置 LoRA：高效微调的关键​

四、数据预处理与整理​

五、训练参数配置​

六、启动蒸馏训练​

七、合并并保存模型​

八、推理测试：看看小模型学到了什么？​

总结​