DistillKit

一个灵活且生产就绪的大语言模型知识蒸馏工具包，支持在线和离线蒸馏工作流，并配备先进的 logits 压缩技术。

DistillKit 为 Arcee 的多款热门开源模型的训练提供了强大支持，包括 Virtuoso、SuperNova Medius 和 Blitz。

特性

• 在线蒸馏：学生模型训练过程中实时进行教师推理
• 离线蒸馏：基于预先捕获的教师输出进行训练，并采用高级压缩技术
• 先进的 Logits 压缩：创新的多项式近似 + 量化 + 位打包技术，在保持蒸馏质量的同时实现极高的压缩比
• 灵活的损失函数：可组合的损失函数，包括 KL 散度、JSD、TVD、排序损失以及隐藏状态对齐等
• 稀疏与稠密支持：高效处理 top-k 稀疏分布或精确的稠密分布
• 经过实战检验：支撑 Arcee 蒸馏模型发布的基础设施
• HuggingFace 集成：基于 Transformers、TRL 和 Accelerate 构建

为什么选择 DistillKit？

虽然在线蒸馏相对简单，但大规模的离线蒸馏需要精心设计。当在数十亿个 token 上进行蒸馏时，仅存储 top-k 的 token-logits 对就会变得极其昂贵。

DistillKit 的压缩系统是数月实验的结果，旨在平衡存储成本、内存吞吐量和蒸馏质量之间的微妙关系。我们的方法如下：

1. 对 logits 分布曲线进行多项式近似
2. 对残差进行误差扩散量化以保持质量
3. 使用任意位宽（1-64 位）进行位级打包

这使得原本不可行的离线蒸馏工作流程得以实际运行。

安装

git clone https://github.com/arcee-ai/distillkit.git
cd distillkit
pip install -e .

可选：Logits 捕获

要捕获您自己的教师输出，请安装捕获依赖项：

pip install -e ".[capture]"

对于大多数用户，我们建议从我们提供的预捕获教师数据集开始（见下文的 数据集）。

快速入门

离线蒸馏

使用预捕获的教师输出训练学生模型：

# config.yaml
project_name: my-distillation
model: Qwen/Qwen3-8B
output_path: ./output
sequence_length: 8192

dataset:
  train_dataset:
    repo_id: arcee-ai/Qwen3-235B-Logits-Packed-8192  # 预捕获的教师输出
    split: train
  prepacked: true

teacher:
  kind: dataset
  logprob_compressor:
    d: 151936  # 词汇表大小
    delta_encoding: true
    error_diffusion: false
    exact_dtype: float32
    exact_k: 32
    k: 128
    polynomial_terms: [0, 1, 2]
    residual_bins: []
    term_dtype: float32

loss_functions:
  - function: cross_entropy
    weight: 0.5
  - function: kl
    weight: 0.5
    temperature: 1.0
    missing_probability_handling: zero
    sparse_chunk_length: 1024

training_args:
  num_train_epochs: 1
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  learning_rate: 2.0e-6
  bf16: true
  optim: adamw_torch
  gradient_checkpointing: true

运行训练：

distillkit config.yaml

在线蒸馏

对于教师在学生训练过程中实时运行的在线蒸馏，完整的配置示例请参阅 examples/afm_test.yml。

核心概念

大语言模型的知识蒸馏

知识蒸馏是将知识从一个（可能更大的）“教师”模型转移到“学生”模型的过程。与仅基于硬标签（正确 token）进行训练不同，学生模型会学习教师对各个 token 的概率分布，这是一种更为丰富的学习信号。

主要优势：

• 更小、更快的模型，同时保持竞争力能
• 更低的推理成本
• 更易于在资源受限的环境中部署

在线与离线蒸馏的区别

在线蒸馏：

• 教师在学生训练过程中实时运行
• 无存储开销
• 适用场景：当您的显存足以同时容纳教师和学生模型，并且可以处理稠密分布时

离线蒸馏：

• 教师输出被预先捕获并压缩
• 可以用同一个教师模型训练多个学生模型
• 适用场景：显存有限、需要重复利用教师信号，或者进行大规模训练时

经验法则： 如果您能够在显存中同时容纳教师和学生的稠密分布，则使用在线蒸馏；否则，使用我们的压缩系统进行离线蒸馏是更好的选择。

稀疏与稠密分布

稠密分布包含整个词汇表的概率。这种方式更准确，但占用大量内存。

稀疏分布只存储 top-k 的 token，是对完整稠密分布的一种有损但实用且高效的近似。在充足的训练数据下，稀疏蒸馏可以达到与稠密蒸馏相当的效果。

DistillKit 同时支持这两种方式，并通过自动分块实现长序列的高效处理。

Logits 压缩

我们的压缩系统在存储效率和蒸馏质量之间取得了平衡：

1. 从教师输出中选择 top-k logits
2. 按对数概率排序，可选应用差分编码
3. 对分布曲线进行多项式拟合
4. 对残差进行量化，可选误差扩散
5. 将所有内容打包成字节向量

您可以调整许多参数，以找到适合您特定需求的存储与保真度之间的权衡。

推荐配置（Arcee 新捕获时使用的配置）：

logprob_compressor:
  d: <your_vocab_size_here>
  k: 128
  exact_k: 16
  exact_dtype: bfloat16
  polynomial_terms: [0, 1, 2, 3, 4, "sqrt"]
  term_dtype: float32
  residual_bins: []
  delta_encoding: false
  error_diffusion: false

这种配置每 token 占用约 300 字节（未压缩分布大小的 0.15%），且质量损失极小。

如果您对存储空间较为紧张，可以尝试以下经济型配置：

logprob_compressor:
  d: <your_vocab_size_here>
  k: 50
  exact_k: 1
  exact_dtype: bfloat16
  polynomial_terms: [0, 1, "sqrt"]
  term_dtype: float32
  residual_bins: []
  delta_encoding: false
  error_diffusion: false

该配置每 token 仅需约 114 字节，体积更小，且重建质量优于直接以 bfloat16 存储 top-32 logits。

请注意，用于捕获 logits 的配置必须与蒸馏配置保持一致。混合使用可能会导致效果不佳。

配置指南

损失函数

DistillKit 支持可组合的损失函数，并为每个损失函数分配独立的权重：

基于分布的损失

• kl：Kullback-Leibler 散度（标准蒸馏损失）
• jsd：Jensen-Shannon 散度（KL 的对称替代方案）
• tvd：总变差距离

排序损失

• hinge：Hinge 排序损失
• logistic_ranking：Logistic 排序损失

隐藏状态对齐

• hs_mse：教师模型与学生模型隐藏状态之间的均方误差
• hs_cosine：隐藏状态之间的余弦相似度

标准损失

• cross_entropy：标准语言建模损失

所有基于分布的损失都支持稀疏和密集两种模式。可以组合使用多种损失函数：

loss_functions:
  - function: cross_entropy
    weight: 0.25
  - function: kl
    weight: 0.5
    temperature: 2.0
  - function: hs_cosine
    weight: 0.25

教师模型配置

离线（从数据集获取）：

teacher:
  kind: dataset
  logprob_compressor:
    d: 128256
    k: 128
    exact_k: 16
    delta_encoding: true
    ...
  # 或：
  legacy_logit_compression:
    vocab_size: 128256
    k: 128
    exact_k: 32
    polynomial_degree: 8
    ...

在线（HuggingFace 模型）：

teacher:
  kind: hf
  path: Qwen/Qwen3-8B
  kwargs: # 加载教师模型时传递的关键字参数
    attn_implementation: flash_attention_2
    torch_dtype: bfloat16

高级主题

压缩深度解析

压缩系统支持两种模式：

传统压缩（完全基于多项式）：

legacy_logit_compression:
  vocab_size: 128256       # 教师词汇表大小
  k: 128                   # 每个 token 的总对数概率数，精确值加近似值
  exact_k: 32              # 以浮点数存储的对数概率数量
  polynomial_degree: 8     # 近似多项式的次数
  with_sqrt_term: false    # 在多项式中包含平方根项
  term_dtype: float32      # 多项式系数的精度
  invert_polynomial: true  # 为改善数值特性而反转多项式

分布量化（较新、更灵活）：

logprob_compressor:
  d: 128256                # 教师词汇表大小
  k: 128                   # 每个 token 的总对数概率数，精确值加近似值
  exact_k: 16              # 以浮点数存储的对数概率数量
  exact_dtype: bfloat16    # “精确”对数概率的数据类型
  delta_encoding: false    # 以差值形式存储对数概率（不推荐）
  error_diffusion: false   # 执行误差扩散以将量化误差分散到各个值上（不推荐）
  polynomial_terms:        # 用于近似尾部的多项式项列表
    - 0
    - 1
    - 2
    - “sqrt”
  term_dtype: float32      # 存储多项式系数的数据类型
  residual_bins:           # 可选的存储量化残差与近似尾部之差的桶列表
    - scale_dtype: float16 # 此桶的比例因子数据类型
      element_bits: 8      # 每个元素的位数
      num_elements: 16     # 此桶中的元素总数
    - scale_dtype: float32 # bfloat16 也可用
      element_bits: 2      # 可以使用最多 64 位
      num_elements: 64
    ...

隐藏状态蒸馏

将学生模型的隐藏状态与教师模型的隐藏状态对齐：

layer_mapping: all  # 或指定层对
loss_functions:
  - function: hs_mse
    weight: 0.5

对于跨架构蒸馏，隐藏状态会通过学习得到的线性映射进行投影。你也可以通过设置 force_hidden_state_projection: true 来为同架构蒸馏启用此功能。

捕获教师输出

要创建自己的离线蒸馏数据集：

python -m distillkit.sample_logits_vllm \
  --model meta-llama/Llama-3.1-70B \
  --dataset allenai/tulu-3-sft-mixture \
  --output ./llama3_70b_tulu_logits/ \
  --compression-config ./compression_config.yaml

需要 vLLM（参见【可选对数概率捕获】）。

内存管理技巧

对于长序列：

• 使用 sparse_chunk_length 将序列分块处理（例如 1024）
• 使用 DeepSpeed ZeRO 第 1 或第 2 阶段来容纳更多 token

一般节省内存的方法：

• 使用 optim: paged_adamw_8bit 或 optim: adamw_bnb_8bit
• 启用 Flash Attention 2：use_flash_attention: true
• 使用 bfloat16 而不是 float32
• 启用 gradient_checkpointing
• 减小批量大小，增加梯度累积步数

示例

• 离线蒸馏（70B → 8B）：examples/llama_70b_base.yml
• 带隐藏状态的在线蒸馏：examples/afm_test.yml
• 多模态模型蒸馏：examples/mistral3.yaml

数据集

我们发布了多个预先捕获的教师数据集：

• Qwen3-235B 指令遵循数据集：约 15 亿 tokens 的通用指令数据，上下文长度为 8192
• DeepSeek V3/R1 合成混合模式推理数据集：约 50 亿 tokens，来自 DeepSeek V3 和 R1，带有前缀以区分推理与非推理轨迹，上下文长度为 16k
• DeepSeek V3 基础模型数据集：约 12 亿 tokens 的原始完成数据，来自 DCLM，由 DeepSeek V3 基础模型捕获

跨架构蒸馏

DistillKit 可以与 mergekit-tokensurgeon 结合使用，实现跨分词器、跨架构的蒸馏。许多 Arcee 模型同时结合了这两种工具：

1. 使用 tokensurgeon 将学生模型的嵌入适配到教师模型的分词器
2. 使用 DistillKit 将教师的知识蒸馏到学生模型
3. 可选地再转换回学生模型原有的分词器，甚至进行其他奇特的合并操作，随心所欲

训练提示

• 从 ~0.5 的交叉熵权重开始，然后根据数据集的质量调整权重
• 蒸馏温度：temperature: 2.0 是一个不错的初始选择
• 缺失概率处理：使用 zero 只关注教师最自信的预测；使用 uniform 则同时匹配教师的不确定性

引用

如果你在研究中使用了 DistillKit，请引用以下内容：

@software{distillkit2024,
  title = {DistillKit：大型语言模型的灵活知识蒸馏},
  author = {Goddard, Charles 和 Atkins, Lucas},
  year = {2024},
  publisher = {Arcee AI},
  url = {https://github.com/arcee-ai/distillkit}
}

社区与支持

• 问题反馈: GitHub Issues
• 讨论交流: Arcee Discord

注意：DistillKit 是一个开源研究项目。虽然它为我们的多款生产级模型提供了支持，并且在资源允许的情况下我们会尽力解决遇到的问题，但社区支持属于尽力而为的范畴。

许可证

DistillKit 采用 Apache License 2.0 许可证发布。

致谢

• Flash Attention 的打包实现改编自 Functionary（MIT 许可证）
• 基于 HuggingFace Transformers、TRL 和 Accelerate 构建

---

由 Arcee AI 竭诚打造

DistillKit 快速上手指南

DistillKit 是一个灵活且生产就绪的大语言模型（LLM）知识蒸馏工具包，支持在线和离线蒸馏工作流，并具备先进的 Logit 压缩技术。它已被用于训练 Arcee 的多个开源模型（如 Virtuoso、SuperNova Medius 等）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• Python: 3.9 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU（建议显存充足以运行教师/学生模型）
• 前置依赖:
  • PyTorch (与 CUDA 版本匹配)
  • Hugging Face Transformers
  • TRL (Transformer Reinforcement Learning)
  • Accelerate
  • (可选) vLLM：如果您需要自行捕获教师模型的输出数据

提示：国内用户建议在 pip 命令中添加国内镜像源（如清华源 -i https://pypi.tuna.tsinghua.edu.cn/simple）以加速依赖下载。

安装步骤

1. 克隆仓库并安装核心包

git clone https://github.com/arcee-ai/distillkit.git
cd distillkit
pip install -e .

2. (可选) 安装日志捕获依赖

如果您计划自行运行教师模型来生成蒸馏数据（而非使用预生成的数据集），需要安装额外的捕获依赖：

pip install -e ".[capture]"

注意：对于大多数初学者，建议直接使用官方提供的预捕获教师数据集（见下文“基本使用”），无需执行此步。

基本使用

DistillKit 的核心优势在于离线蒸馏，即利用预先生成并压缩的教师模型输出数据进行训练，从而节省显存和计算资源。

场景：使用预捕获数据进行离线蒸馏

以下示例展示如何配置并启动一个从大型教师模型（如 Qwen3-235B）到小型学生模型（如 Qwen3-8B）的蒸馏任务。

1. 创建配置文件

新建一个名为 config.yaml 的文件，填入以下内容：

# config.yaml
project_name: my-distillation
model: Qwen/Qwen3-8B  # 学生模型
output_path: ./output
sequence_length: 8192

dataset:
  train_dataset:
    repo_id: arcee-ai/Qwen3-235B-Logits-Packed-8192  # 使用官方预捕获的教师 Logits 数据集
    split: train
  prepacked: true

teacher:
  kind: dataset  # 指定为离线数据集模式
  logprob_compressor:
    d: 151936  # 词表大小 (需与教师模型一致)
    delta_encoding: true
    error_diffusion: false
    exact_dtype: float32
    exact_k: 32
    k: 128
    polynomial_terms: [0, 1, 2]
    residual_bins: []
    term_dtype: float32

loss_functions:
  - function: cross_entropy
    weight: 0.5
  - function: kl
    weight: 0.5
    temperature: 1.0
    missing_probability_handling: zero
    sparse_chunk_length: 1024

training_args:
  num_train_epochs: 1
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  learning_rate: 2.0e-6
  bf16: true
  optim: adamw_torch
  gradient_checkpointing: true

2. 启动训练

在终端中运行以下命令开始蒸馏：

distillkit config.yaml

关键配置说明

• dataset.repo_id: 指向 Hugging Face 上的预压缩 Logits 数据集。使用这些数据可以避免在训练时加载巨大的教师模型。
• teacher.kind: 设置为 dataset 表示启用离线模式；若需在线蒸馏（实时推理教师模型），则设为 hf 并指定模型路径。
• loss_functions: 支持组合多种损失函数（如交叉熵 cross_entropy 和 KL 散度 kl），通过 weight 调整权重。
• logprob_compressor: 必须与生成数据集时使用的压缩配置保持一致，否则无法正确解码 Logits。

进阶提示

• 在线蒸馏: 如果显存足够同时容纳教师和学生模型，可参考 examples/afm_test.yml 配置在线蒸馏，无需预先存储 Logits。
• 显存优化: 对于长序列训练，建议在配置中启用 gradient_checkpointing: true 并使用 sparse_chunk_length 分块处理，或结合 DeepSpeed ZeRO 技术。
• 自定义数据捕获: 若需针对特定数据集生成教师 Logits，可使用 python -m distillkit.sample_logits_vllm 命令（需先安装 vLLM）。