利用一张GPU卡动手实践：从零开始蒸馏小模型

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

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

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

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

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

如需 A100 卡资源做验证，可注册 Spader.AI 在线使用 GPU 算力资源。

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

我们使用由 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