蒸馏 Trainer

概述

Distillation Trainer 实现了 Rishabh Agarwal、Nino Vieillard、Yongchao Zhou、Piotr Stanczyk、Sabela Ramos、Matthieu Geist 和 Olivier Bachem 在 On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes 中描述的在线策略知识蒸馏。

该论文的摘要如下：

知识蒸馏（KD）被广泛用于通过训练一个更小的学生模型来压缩教师模型，从而降低其推理成本和内存占用。然而，当前针对自回归序列模型的知识蒸馏方法存在训练期间看到的输出序列与学生模型在推理期间生成的序列之间的分布不匹配问题。为了解决这一问题，我们提出了广义知识蒸馏（GKD）。GKD 不再仅仅依赖固定的输出序列集合，而是利用教师模型对这些序列的反馈，在学生模型自生成的输出序列上训练学生模型。与有监督的知识蒸馏方法不同，GKD 还提供了在学生模型和教师模型之间使用替代损失函数的灵活性，当学生模型缺乏模仿教师模型分布的表达力时，这一点尤为有用。

DistillationTrainer 训练一个较小的学生模型，使其在学生自身的 on-policy 生成上匹配教师的下一个 token 分布。它会在 on-policy 模式下生成学生的补全（可选由 vLLM 提供支持），并使用内存高效的分块 Jensen-Shannon 散度损失，在教师的完整下一个 token 分布上进行匹配，从而避免在内存中完整实例化教师的稠密分布。

此训练器由 Carlos Miguel Patiño 贡献。

对于一种将生成与训练解耦、并通过 HTTP 提供的教师模型而非本地前向传播进行评分的异步对应方案，请参阅 AsyncDistillationTrainer。

快速入门

本示例演示了如何使用蒸馏方法训练模型。我们使用UltraFeedback 提示数据集中的提示，从Qwen 2.5 1.5B Instruct 教师模型中蒸馏出一个Qwen 2.5 0.5B Instruct 模型。您可以在这里查看该数据集中的数据：

以下是用于训练模型的脚本。

# train_distillation.py
from datasets import load_dataset
from trl import DistillationTrainer

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

使用以下命令执行该脚本：

torchrun --nproc_per_node 8 train_distillation.py

深入了解蒸馏方法

On-policy 知识蒸馏训练学生模型在学生自身生成的补全上复现教师模型的下一个 token 分布，而不是在固定的教师输出集合上进行。从自身生成中学习使学生能够纠正自身的错误，这通常优于 off-policy 蒸馏。本节将详细解析其实际运作方式，涵盖两个关键步骤：生成补全和计算损失。

生成补全

在每个训练步骤中，学生模型会为采样的提示生成一批补全内容。

计算损失

损失函数是 广义 Jensen-Shannon 散度（JSD） ，计算于学生分布 $p_S$ 与教师分布 $p_T$ 在生成的补全 token 上，通过 beta 进行插值，定义如下： $\mathcal{L}_\beta = \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_M \right] + (1 - \beta) \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_M \right], \qquad p_M = (1 - \beta) \, p_S + \beta \, p_T,$

其中 $p_M$ 是这两个分布的 $\beta$ -混合。端点退化为纯散度： beta=0.0 给出前向 KL $\mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_S \right]$ 而 beta=1.0 给出反向 KL $\mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_T \right]$ .

在实践中，投影到词汇表 logits 以及散度的计算是分块进行的，因此峰值激活内存不会随完整的词汇表 × 序列长度 logits 张量而增长。请参阅减少内存使用。

预期的数据集类型

数据集应格式化为对话式的仅提示数据集。由于学生模型在策略内自行生成补全，因此只需要提示：

{"prompt": [{"role": "user", "content": "What color is the sky?"}]}

记录的指标

在训练和评估过程中，我们记录以下指标：

num_tokens：迄今为止处理的 token 总数，包括 prompt 和 completion。使用工具时，仅统计非工具 token。
step_time：每个训练步骤（包括生成）的平均耗时（以秒为单位）。
completions/mean_length：生成补全的平均长度。使用工具时，仅统计非工具 token。
completions/min_length：生成补全的最小长度。使用工具时，仅统计非工具 token。
completions/max_length：生成补全的最大长度。使用工具时，仅计算非工具 token。
completions/mean_terminated_length：以 EOS 终止的生成补全的平均长度。使用工具时，仅统计非工具 token。
completions/min_terminated_length：以 EOS 终止的生成补全的最小长度。使用工具时，仅统计非工具 token。
completions/max_terminated_length：以 EOS 终止的生成补全的最大长度。使用工具时，仅统计非工具 token。
completions/clipped_ratio：被截断（裁剪）的补全比例。
tools/call_frequency：生成批次中每次补全的平均工具调用次数。仅在提供 tools 时记录。
tools/failure_frequency：工具调用失败的比例（工具未找到、抛出异常或调用类型不受支持）。当未调用任何工具时，该值为 0.0。仅在提供 tools 时记录。
entropy：生成的补全中 token 预测的平均熵（以 nats 为单位）。

自定义

使用 vLLM 驱动的生成来加速训练

在使用 on-policy 方法训练时，生成过程往往是主要的性能瓶颈。为了加速生成，你可以使用 vLLM，这是一个面向 LLM 的高吞吐量、低延迟推理引擎。要启用它，首先使用以下命令安装该包

pip install trl[vllm]

我们支持在训练期间使用 vLLM 的两种方式：colocate 模式和server 模式。

选项 1：Colocate 模式

在此模式下，vLLM 在训练器进程内运行，并与训练模型共享 GPU 内存。这避免了启动单独的服务器，并可能提高 GPU 利用率，但可能导致训练 GPU 上的内存争用。这是默认模式。

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

选项 2：服务器模式

在此模式下，vLLM 在独立进程中运行（并使用独立的 GPU），并通过 HTTP 与训练器进行通信。如果您拥有专门用于推理的 GPU，则此模式非常理想。

启动 vLLM 服务器：

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

在训练脚本中启用服务器模式：

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

请确保服务器使用的 GPU 与训练器不同，否则可能会遇到 NCCL 错误。您可以通过 CUDA_VISIBLE_DEVICES 环境变量来指定要使用的 GPU。

根据模型大小以及训练所需的整体 GPU 内存，您可能需要调整 DistillationConfig 中的 vllm_gpu_memory_utilization 参数，以避免资源利用不足或内存溢出错误。

有关更多信息，请参阅使用 vLLM 加速训练。

使用 PEFT 训练适配器

我们支持与 🤗 PEFT 库的紧密集成，让您能够训练适配器并将其共享到 Hub，而无需训练整个学生模型。

from datasets import load_dataset
from trl import DistillationTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)
trainer.train()

蒸馏损失会直接读取 lm_head.weight，并在不经过 PeftModel.forward() 的情况下运行学生骨干网络。因此，针对 lm_head 的适配器（通过 target_modules 指定）以及提示学习方法（PromptTuning、PrefixTuning、P-Tuning）会被拒绝，因为它们会被静默忽略。若要训练该头部，请改用 modules_to_save=["lm_head"]。

Agent 训练

DistillationTrainer 支持智能体训练：学生模型在生成过程中调用工具，并基于由此产生的轨迹进行蒸馏。工具返回结果的 token 会在损失计算中被屏蔽，因此学生模型仅针对其自身生成的 token 进行训练。

工具

tools 参数需要一个 Python 函数列表，用于定义智能体可用的工具：

from trl import DistillationTrainer

trainer = DistillationTrainer(
    tools=[tool1, tool2],
    ...,
)

每个工具必须是一个标准的 Python 函数，具有带类型提示的参数和返回类型，并附带一个Google 风格的 docstring，用于描述其用途、参数和返回值。有关更多详细信息，请参阅传递工具指南。

工具调用循环要求聊天模板是前缀保持的（追加一条工具消息时，不得改变之前消息的渲染方式）。对于已知的模型系列（例如 Qwen3、DeepSeek-V3），当启用工具时，TRL 会自动替换为经过修补的训练模板。完整的列表请参阅聊天模板。

在 DistillationConfig 中使用 max_tool_calling_iterations 来限制工具调用轮次的数量。默认情况下没有上限，当学生模型生成一个不包含工具调用的响应轮次、当补全内容（所有轮次和工具结果）达到 max_completion_length 个 token，或者当完整序列达到模型的最大长度（在 vllm_mode="colocate" 时为 vllm_max_model_length，否则为模型的 max_position_embeddings）时，生成将停止。

示例：

from trl import DistillationTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

trainer = DistillationTrainer(
    tools=[multiply],
    ...,
)

暂不支持异步工具——请传入同步函数。

多模态工具响应

工具可以通过返回内容块列表，在文本之外同时返回图像。这在 VLM 智能体训练中非常有用，因为工具可以提供视觉反馈（例如，屏幕截图、图表、摄像头捕获）。

from PIL import Image

def take_screenshot() -> list:
    """
    Takes a screenshot of the current screen.

    Returns:
        The screenshot image with a description.
    """
    img = Image.open("screenshot.png")
    return [{"type": "image", "image": img}, {"type": "text", "text": "Here is the screenshot."}]

返回的图像会自动注入到对话中，并传递给 VLM 用于后续的生成轮次。

训练视觉语言模型

DistillationTrainer 支持在包含文本和图像的多模态数据集上蒸馏视觉语言模型（VLM）。将 VLM 同时作为学生模型和教师模型传入，并提供一个仅包含提示词的数据集，该数据集需包含 image 列（每个样本一张图像）或 images 列（每个样本一个图像列表）。有关预期数据集结构的更多信息，请参阅 Dataset Format — Vision datasets 部分。

已测试：

Gemma 3 — 例如 google/gemma-3-4b-it
LLaVA-NeXT — 例如 llava-hf/llava-v1.6-mistral-7b-hf
Qwen2-VL — 例如 Qwen/Qwen2-VL-2B-Instruct
Qwen2.5-VL — 例如 Qwen/Qwen2.5-VL-3B-Instruct

无法保证与所有 VLM 的兼容性。如果您认为某个模型应当被支持，欢迎在 GitHub 上提交 issue——或者更好的是，提交一个包含所需更改的 pull request。

命令行界面

使用 trl distillation CLI 从命令行启动蒸馏训练。它通过标准的 ModelConfig 标志支持完整训练和 LoRA。

# Full training:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-5 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1

# LoRA:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-4 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16

DistillationTrainer

类 trl.DistillationTrainer

< source >

( model: str | PreTrainedModel | PeftModelteacher_model: str | transformers.modeling_utils.PreTrainedModel | None = Noneargs: trl.trainer.distillation_config.DistillationConfig | None = Nonetrain_dataset: datasets.arrow_dataset.Dataset | None = Noneeval_dataset: datasets.arrow_dataset.Dataset | dict[str, datasets.arrow_dataset.Dataset] | None = Noneprocessing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[typing.Optional[torch.optim.Optimizer], typing.Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None)quantization_config: BitsAndBytesConfig | None = Nonepeft_config: PeftConfig | None = Nonetools: list[collections.abc.Callable] | None = None )

参数

model (str 或 PreTrainedModel 或 PeftModel) — 要训练的模型。可以是以下之一：
- 一个字符串，可以是托管在 huggingface.co 模型仓库中的预训练模型的 model id，也可以是包含使用 save_pretrained 保存的模型权重的目录的路径，例如 './my_model_directory/'。该模型使用 <ModelArchitecture>.from_pretrained 加载（其中 <ModelArchitecture> 由模型配置推导得出），并传入 args.model_init_kwargs 中的关键字参数。如果未在 args.model_init_kwargs 中指定 dtype，则默认为 float32。这与 from_pretrained 不同，在后者中（自 Transformers v5 起），dtype 是从模型配置中推断得出的。
- 一个 PreTrainedModel 对象。仅支持因果语言模型。
- PeftModel 对象。仅支持因果语言模型。
teacher_model (str 或 PreTrainedModel，可选) — 学生模型训练以匹配其下一个 token 分布的教师模型。可以是 模型 ID / 路径（加载方式与 model 相同，使用 args.teacher_model_init_kwargs）或一个已实例化的 PreTrainedModel。它必须与学生模型共享词汇表。子类可以省略此参数，通过其他方式提供教师模型（例如远程服务器）。
args (DistillationConfig，可选) — 此训练器的配置。如果为 None，则使用默认配置。
train_dataset (Dataset 或 IterableDataset) — 用于训练的 Dataset。它必须包含一个名为 "prompt". 数据集中的其他列将被忽略。样本的格式可以是以下两种之一：
- 标准：每个样本包含纯文本。
- 对话式：每个样本包含结构化的消息（例如角色和内容）。
当 train_dataset 是一个 IterableDataset（例如流式数据集）时，必须在训练参数中设置 max_steps，因为其长度无法推断，而确定训练循环的上限并配置学习率调度器需要知道总的训练步数。
eval_dataset (Dataset、IterableDataset、DatasetDict、IterableDatasetDict 或 dict[str, Dataset | IterableDataset]) — 用于评估的数据集。它必须满足与 train_dataset 相同的要求。
processing_class (PreTrainedTokenizerBase、ProcessorMixin、可选) — 用于处理数据的处理类。填充侧必须设置为“left”。如果为 None，则使用 from_pretrained 从模型名称加载处理类。必须设置填充标记 tokenizer.pad_token。如果处理类未设置填充标记，则默认使用 tokenizer.eos_token。
callbacks (列表 TrainerCallback, 可选) — 用于自定义训练循环的回调列表。这些回调将被添加到此处.
如果要移除其中一个默认使用的回调，请使用 remove_callback 方法。
optimizers (tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]，可选，默认为 (None, None)) — 一个包含要使用的优化器和调度器的元组。默认情况下，将在您的模型上使用 AdamW 的实例，并使用由 args 控制的、通过 get_linear_schedule_with_warmup 给出的调度器。
quantization_config (BitsAndBytesConfig, 可选) — 从模型标识符加载模型时使用的量化配置。与 peft_config 结合用于 QLoRA 训练。如果模型已经实例化，则忽略此配置。
peft_config (PeftConfig，可选) — 用于包装模型的 PEFT 配置。如果为 None，则不包装模型。
tools（Callable 列表，可选）— 一个可调用工具函数的列表，模型可在生成过程中调用这些工具。每个工具应是一个标准的 Python 函数，具有正确类型提示的参数和返回值，以及一个 Google 风格的 docstring，用于描述其用途、参数和返回值。更多详情，请参阅： https://huggingface.co/docs/transformers/en/chat_extras#passing-tools。模型利用函数的名称、类型提示和 docstring 来确定如何调用它。请确保模型的 chat template 支持工具使用，并且该模型已针对工具调用进行了微调。

用于知识蒸馏的 Trainer。学生模型采用 on-policy 方式进行训练——即由学生模型自行生成补全内容——以在广义 Jensen-Shannon 散度下匹配教师模型的下一个 token 分布（通过 beta 对前向 KL、反向 KL 和 JSD 进行插值），该方法在 On-Policy Distillation of Language Models 中提出。

示例：

>>> from trl import DistillationTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/tldr", split="train")

>>> trainer = DistillationTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

训练

< source >

参数

resume_from_checkpoint (str 或 bool，可选) — 如果为 str，则为之前 Trainer 实例保存的检查点的本地路径。如果为 bool 且等于 True，则加载之前 Trainer 实例保存在 args.output_dir 中的最后一个检查点。如果存在，训练将从此处加载的模型/优化器/调度器状态中恢复。
trial (optuna.Trial 或 dict[str, Any]，可选) — 用于超参数搜索的试验运行或超参数字典。
ignore_keys_for_eval (list[str]，可选) — 在训练期间收集预测结果以进行评估时，应忽略的模型输出（如果输出是字典）中的键列表。

返回值

~trainer_utils.TrainOutput

包含全局步数、训练损失和指标的 Object。

主要训练入口。

save_model

< source >

( output_dir: str | None = None_internal_call: bool = False )

将保存模型，以便您可以使用 from_pretrained() 重新加载它。

仅从主进程保存。

push_to_hub

< source >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = Noneprivate: bool | None = None**kwargs )

参数

commit_message (str，可选，默认为 "End of training") — 推送时提交的提交信息。
blocking (bool，可选，默认为 True) — 是否仅在 git push 完成后才返回。
token (str，可选，默认为 None) — 具有写入权限的 Token，用于覆盖 Trainer 的原始参数。
revision (str, 可选) — 要从中提交内容的 git 版本。默认为 “main” 分支的头部。
private (bool，可选) — 创建仓库时是否将其设为私有。默认值为 args.hub_private_repo。如果仓库已存在，则此参数无效。
kwargs (dict[str, Any]，可选) — 传递给 ~Trainer.create_model_card 的额外关键字参数。

将 self.model 和 self.processing_class 上传到 🤗 模型库中的仓库 self.args.hub_model_id。

DistillationConfig

类 trl.DistillationConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 1e-06lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear'lr_scheduler_kwargs: dict | str | None = Nonewarmup_steps: float = 0optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused'optim_args: str | None = Noneweight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08optim_target_modules: None | str | list[str] = Nonegradient_accumulation_steps: int = 1average_tokens_across_devices: bool = Truemax_grad_norm: float = 1.0label_smoothing_factor: float = 0.0bf16: bool | None = Nonefp16: bool = Falsebf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: bool | None = Noneactivation_checkpointing: bool | None = Noneactivation_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonegradient_checkpointing: bool = Truegradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonetorch_compile: bool = Falsetorch_compile_backend: str | None = Nonetorch_compile_mode: str | None = Noneuse_liger_kernel: bool = Falseliger_kernel_config: dict[str, bool] | None = Noneuse_cache: bool = Falseneftune_noise_alpha: float | None = Nonetorch_empty_cache_steps: int | None = Noneauto_find_batch_size: bool = Falselogging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps'logging_steps: float = 10logging_first_step: bool = Falselog_on_each_node: bool = Truelogging_nan_inf_filter: bool = Trueinclude_num_input_tokens_seen: str | bool = 'no'log_level: str = 'passive'log_level_replica: str = 'warning'disable_tqdm: bool | None = Nonereport_to: None | str | list[str] = 'none'run_name: str | None = Noneproject: str = 'huggingface'trackio_space_id: str | None = Nonetrackio_bucket_id: str | None = Nonetrackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = Noneeval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no'eval_steps: float | None = Noneeval_delay: float = 0per_device_eval_batch_size: int = 8prediction_loss_only: bool = Falseeval_on_start: bool = Falseeval_do_concat_batches: bool = Trueeval_use_gather_object: bool = Falseeval_accumulation_steps: int | None = Noneinclude_for_metrics: list[str] = <factory>batch_eval_metrics: bool = Falsesave_only_model: bool = Falsesave_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps'save_steps: float = 500save_on_each_node: bool = Falsesave_total_limit: int | None = Noneenable_jit_checkpoint: bool = Falsepush_to_hub: bool = Falsehub_token: str | None = Nonehub_private_repo: bool | None = Nonehub_model_id: str | None = Nonehub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save'hub_always_push: bool = Falsehub_revision: str | None = Noneload_best_model_at_end: bool = Falsemetric_for_best_model: str | None = Nonegreater_is_better: bool | None = Noneignore_data_skip: bool = Falserestore_callback_states_from_checkpoint: bool = Falsefull_determinism: bool = Falseseed: int = 42data_seed: int | None = Noneuse_cpu: bool = Falseaccelerator_config: dict | str | None = Noneparallelism_config: accelerate.parallelism_config.ParallelismConfig | None = Nonedataloader_drop_last: bool = Falsedataloader_num_workers: int = 0dataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falsedataloader_prefetch_factor: int | None = Nonedataloader_multiprocessing_context: str | None = Nonedataloader_in_order: bool = Trueremove_unused_columns: bool | None = Falselabel_names: list[str] | None = Nonetrain_sampling_strategy: str = 'random'length_column_name: str = 'length'ddp_find_unused_parameters: bool | None = Noneddp_bucket_cap_mb: int | None = Noneddp_broadcast_buffers: bool | None = Noneddp_static_graph: bool | None = Noneddp_backend: str | None = Noneddp_timeout: int = 1800fsdp: str | None = Nonefsdp_config: dict[str, typing.Any] | str | None = Nonedeepspeed: dict | str | None = Nonedebug: str | list[transformers.debug_utils.DebugOption] = ''skip_memory_metrics: bool = Truedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseresume_from_checkpoint: str | None = Nonelocal_rank: int = -1model_init_kwargs: dict[str, typing.Any] | str | None = Nonetrust_remote_code: bool = Falseteacher_model_name_or_path: str | None = Noneteacher_model_revision: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Nonedisable_dropout: bool = Falsemax_completion_length: int | None = 512ds3_gather_for_generation: bool = Trueshuffle_dataset: bool | None = Truepad_to_multiple_of: int | str | None = 'auto'temperature: float = 1.0top_p: float = 1.0top_k: int = 0min_p: float | None = Nonegeneration_kwargs: dict | str | None = Nonechat_template_kwargs: dict | str | None = Nonerepetition_penalty: float = 1.0cache_implementation: str | None = Noneuse_vllm: bool = Falsevllm_mode: str = 'colocate'vllm_model_impl: str = 'vllm'vllm_enable_sleep_mode: bool = Falsevllm_structured_outputs_regex: str | None = Nonevllm_server_base_url: str | None = Nonevllm_server_host: str = '0.0.0.0'vllm_server_port: int = 8000vllm_server_timeout: float = 240.0vllm_group_port: int = 51216vllm_gpu_memory_utilization: float = 0.3vllm_max_model_length: int | None = Nonevllm_tensor_parallel_size: int = 1beta: float = 1.0max_tool_calling_iterations: int | None = Nonelog_completions: bool = Falsenum_completions_to_print: int | None = Nonelog_unique_prompts: bool = False )

控制模型和教师模型的参数

model_init_kwargs (str 或 dict[str, Any]，可选) — AutoModelForCausalLM.from_pretrained 的关键字参数，当训练器的 model 参数以字符串形式提供时使用。加载处理类时也会使用 revision 值。
trust_remote_code (bool, optional, defaults to False) — Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to from_pretrained and from_pretrained, for both the student and teacher.
teacher_model_name_or_path (str, 可选) — 教师模型的名称或路径。在本地加载教师模型时使用。
teacher_model_revision (str, 可选) — 教师模型的版本（例如分支名称、标签或提交哈希值）。
teacher_model_init_kwargs (str 或 dict[str, Any]，可选) — 从字符串实例化教师模型时，传递给 AutoModelForCausalLM.from_pretrained 的关键字参数。
disable_dropout (bool，可选，默认为 False) — 是否在训练期间禁用学生模型中的 dropout。

控制数据预处理的参数

remove_unused_columns (bool，可选，默认为 False) — 是否仅保留数据集中的 "prompt" 列。Trainer 会消费原始 prompt 列并按策略生成补全，因此默认值为 False。
max_completion_length (int 或 None，可选，默认为 512) — 在 on-policy 生成期间，每次补全生成的最大 token 数。
ds3_gather_for_generation (bool，可选，默认为 True) — 此设置适用于 DeepSpeed ZeRO-3。如果启用，策略模型权重将在生成时进行聚合，从而提升生成速度。但是，禁用此选项允许训练超出单个 GPU 显存容量的模型，代价是生成速度变慢。禁用此选项与 vLLM 生成不兼容。
shuffle_dataset (bool，可选，默认为 True) — 是否打乱训练数据集。
pad_to_multiple_of (int 或 str，可选，默认为 "auto") — 如果设置，提示词 ID 和补全 ID 将被填充为该值的倍数。"auto" 在 Neuron (AWS Trainium) 上填充为 128 的倍数，因为 Neuron 针对每种输入形状只编译一次，而在其他平台上不进行填充。

控制生成的参数

temperature (float，可选，默认为 1.0) — 用于生成期间采样以及计算蒸馏损失的温度。较高的值会产生更平滑的概率分布。
top_p (float，可选，默认为 1.0) — 用于 on-policy 生成的 Top-p（核）采样参数。
top_k (int，可选，默认为 0) — 用于 on-policy 生成的 Top-k 采样参数。0 表示禁用 top-k 过滤。
min_p (float, 可选) — 最小 token 概率，该值会按最可能 token 的概率进行缩放。它必须是 0.0 和 1.0 之间的值。典型值在 0.01-0.2 范围内。
generation_kwargs (dict[str, Any], 可选) — 在采样补全时传递给 GenerationConfig（如果使用 transformers）或 SamplingParams（如果使用 vLLM）的额外关键字参数。可用于进一步自定义生成行为，例如设置 suppress_tokens、num_beams 等。如果其中包含与其他生成参数（如 min_p、top_p 等）冲突的键，则会覆盖它们。
chat_template_kwargs (dict[str, Any], 可选) — 在生成补全时传递给 apply_chat_template 函数的其他关键字参数。
repetition_penalty (float，可选，默认为 1.0) — 一个浮点数，根据新 token 是否出现在提示词和迄今为止生成的文本中来对其进行惩罚。值 > 1.0 会鼓励模型使用新的 token，而值 < 1.0 会鼓励模型重复使用 token。
cache_implementation (str, 可选) — 当 use_vllm 设置为 False 时，用于加速生成的缓存方法实现。

控制由 vLLM 驱动的生成加速的参数

use_vllm (bool，可选，默认为 False) — 是否使用 vLLM 从学生模型生成 on-policy 补全。
vllm_mode (str，可选，默认为 "colocate") — 学生 vLLM 集成的模式。可以是 "server" 或 "colocate"。
vllm_model_impl (str，可选，默认为 "vllm") — vLLM 的模型实现后端。请使用 "vllm" 或 "transformers"。
vllm_enable_sleep_mode (bool，可选，默认为 False) — 启用 vLLM 睡眠模式，以便在优化器步骤期间卸载学生模型权重。
vllm_structured_outputs_regex (str, 可选) — 用于 vLLM 结构化输出的正则表达式模式。

用于控制 vLLM 服务器的参数（仅在 `vllm_mode` 为 `"server"` 时使用）

vllm_server_base_url (str，可选) — 学生 vLLM 服务器的基础 URL。如果提供，则忽略 vllm_server_host 和 vllm_server_port。
vllm_server_host (str，可选，默认为 "0.0.0.0") — 学生 vLLM 服务器的主机地址。
vllm_server_port (int，可选，默认为 8000) — 学生 vLLM 服务器的端口。
vllm_server_timeout (float，可选，默认为 240.0) — 连接学生 vLLM 服务器的超时时间。
vllm_group_port (int, 可选, 默认为 51216) — vLLM 权重更新组（NCCL 通信器）使用的端口。

控制共置 vLLM 执行的参数（仅在 `vllm_mode` 为 `"colocate"` 时使用）

vllm_gpu_memory_utilization (float，可选，默认值为 0.3) — 用于共置的 student vLLM 引擎的 GPU 内存利用率。
vllm_max_model_length (int, 可选) — 共置 vLLM 引擎的最大模型序列长度。
vllm_tensor_parallel_size (int，可选，默认为 1) — 用于共置学生 vLLM 引擎的张量并行大小。

控制训练的参数

beta (float，可选，默认为 1.0) — 广义 Jensen-Shannon 散度损失的插值系数。当 0.0 时，损失为前向 KL 散度。当 1.0 时，损失为反向 KL 散度。当 0.5 时，为标准 JSD。与 GRPO 的 beta（针对参考模型的 KL 惩罚系数）不同，此处用于选择散度本身；不存在参考模型的 KL 惩罚。
max_tool_calling_iterations (int, 可选) — 训练智能体时，工具调用轮次的最大数量。如果为 None，则没有上限，生成将在以下情况停止：模型生成了一个不包含工具调用的响应轮次；补全内容（所有轮次和工具结果）达到 max_completion_length 个 token；或者完整序列达到模型的最大长度（在 vllm_mode="colocate" 时为 vllm_max_model_length，否则为模型的 max_position_embeddings）。

控制日志记录的参数

log_completions (bool，可选，默认为 False) — 是否每隔 logging_steps 步记录一组 (prompt, completion) 对的样本。如果已安装 rich，则会打印该样本。如果启用了 wandb 和/或 trackio 日志记录，则会将其记录到 wandb 和/或 trackio。
num_completions_to_print (int, 可选) — 使用 rich 打印的补全数量。如果为 None，则记录所有补全。
log_unique_prompts (bool，可选，默认为 False) — 是否记录唯一的提示词。如果为 True，则仅记录唯一的提示词。如果为 False，则记录所有提示词。

DistillationTrainer 的配置类。

扩展了 TrainingArguments，并添加了知识蒸馏特有的参数。所有必需的字段均在此处声明。

使用 HfArgumentParser，我们可以将该类转换为 argparse 参数，从而在命令行中指定。

这些参数的默认值与 TrainingArguments 不同：

logging_steps：默认值为 10，而非 500。

gradient_checkpointing：默认值为 True，而非 False。

bf16：如果未设置 fp16，则默认值为 True，而不是 False。

learning_rate：默认值为 1e-6，而非 5e-5。

Update on GitHub