广义知识蒸馏训练器

概述

广义知识蒸馏（GKD）由 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 还提供了在学生模型和教师模型之间使用替代损失函数的灵活性，当学生模型缺乏模仿教师模型分布的表达力时，这一点尤为有用。此外，GKD 促进了蒸馏与强化学习微调（RLHF）的无缝集成。我们展示了 GKD 在摘要、翻译和算术推理任务上蒸馏自回归语言模型的有效性，以及用于指令微调的任务无关蒸馏。

GKD 的关键方面包括：

它通过在自生成的输出序列上训练学生模型，解决了自回归序列模型中训练与推理之间的分布不匹配问题。
GKD 允许通过广义 Jensen-Shannon 散度（JSD）灵活选择学生模型与教师模型之间的不同散度度量，当学生模型缺乏完全模仿教师模型的能力时，这一点尤为有用。

此训练后方法由 Kashif Rasul 和 Lewis Tunstall 贡献。

使用技巧

experimental.gkd.GKDTrainer 是 SFTTrainer 类的封装，它接受一个教师模型参数。需要通过 experimental.gkd.GKDConfig 设置三个参数，分别是：

lmbda：控制学生数据比例，即 on-policy 学生生成输出的占比。当 lmbda=0.0 时，损失退化为监督式 JSD，学生使用教师的 token 级概率进行训练。当 lmbda=1.0 时，损失退化为 on-policy JSD，学生生成输出序列，并由教师针对这些序列提供 token 级反馈。对于 [0, 1] 之间的值，每个 batch 会根据 lmbda 的值在两种模式之间随机选择。
seq_kd：控制是否执行序列级 KD（可视为在教师生成的输出上进行的监督式 FT）。当 seq_kd=True 且 lmbda=0.0 时，损失函数简化为监督式 JSD，其中教师生成输出序列，学生从教师处获得针对这些序列的 token 级反馈。
beta：控制广义 Jensen-Shannon 散度中的插值。当 beta=0.0 时，损失近似于前向 KL 散度；当 beta=1.0 时，损失近似于反向 KL 散度。对于 [0, 1] 之间的值，它在这两者之间进行插值。

作者发现，on-policy 数据（较高的 lmbda）表现更好，且最优的 beta 值会因任务和评估方法的不同而有所变化。

在训练 Gemma 模型时，请确保 attn_implementation="kernels-community/flash-attn2"。否则，由于该架构采用了 soft capping 技术，您将在 logits 中遇到 NaN。

基本 API 如下：

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl.experimental.gkd import GKDConfig, GKDTrainer

NUM_DUMMY_SAMPLES = 100

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
# The model to optimise
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
# The teacher model to calculate the KL divergence against
teacher_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-1.5B-Instruct")

train_dataset = Dataset.from_dict(
    {
        "messages": [
            [
                {"role": "user", "content": "Hi, how are you?"},
                {"role": "assistant", "content": "I'm great thanks"},
            ]
        ]
        * NUM_DUMMY_SAMPLES
    }
)
eval_dataset = Dataset.from_dict(
    {
        "messages": [
            [
                {"role": "user", "content": "What colour is the sky?"},
                {"role": "assistant", "content": "The sky is blue"},
            ]
        ]
        * NUM_DUMMY_SAMPLES
    }
)

training_args = GKDConfig(output_dir="gkd-model", per_device_train_batch_size=1)
trainer = GKDTrainer(
    model=model,
    teacher_model=teacher_model,
    args=training_args,
    processing_class=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

预期的数据集类型

数据集应格式化为一个“messages”列表，其中每条消息是一个字典列表，包含以下键：

role：可以是 system、assistant 或 user
content：消息内容

GKDTrainer

类 trl.experimental.gkd.GKDTrainer

< source >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str, NoneType] = Noneteacher_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str] = Noneargs: trl.experimental.gkd.gkd_config.GKDConfig | None = Nonedata_collator: collections.abc.Callable[[list[typing.Any]], dict[str, typing.Any]] | 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.image_processing_utils.BaseImageProcessor | transformers.feature_extraction_utils.FeatureExtractionMixin | transformers.processing_utils.ProcessorMixin | None = Nonecompute_metrics: collections.abc.Callable[[transformers.trainer_utils.EvalPrediction], dict] | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)preprocess_logits_for_metrics: collections.abc.Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = Nonepeft_config: PeftConfig | None = Noneformatting_func: collections.abc.Callable | None = None )

参数

model (PreTrainedModel 或 torch.nn.Module 或 str，可选) — 要训练的模型，或从预训练模型实例化模型的字符串标识符。
teacher_model (PreTrainedModel 或 torch.nn.Module 或 str，可选) — 用于知识蒸馏的教师模型，或从预训练模型实例化模型的字符串标识符。
args (experimental.gkd.GKDConfig，可选) — 训练参数。
data_collator (DataCollator, 可选) — 用于将数据集中的样本进行批处理的 Data collator。默认使用 processing_class 的 experimental.utils.DataCollatorForChatML。
train_dataset (Dataset, 可选) — 用于训练的 Dataset。
eval_dataset (Dataset 或 dict 类型的 Dataset，可选) — 用于评估的数据集。
processing_class（PreTrainedTokenizerBase、BaseImageProcessor、FeatureExtractionMixin 或 ProcessorMixin，可选）— 用于处理数据的类。
compute_metrics (Callable, 可选) — 用于在评估时计算指标的函数。必须接收一个 EvalPrediction 并返回一个从字符串到浮点数的字典。
callbacks（list，元素为 TrainerCallback，可选）— 训练期间使用的回调。
optimizers（tuple，包含 torch.optim.Optimizer 和 torch.optim.lr_scheduler.LambdaLR，可选，默认为 (None, None)）— 包含训练所使用的优化器和学习率调度器的元组。
preprocess_logits_for_metrics (Callable，可选) — 用于在计算指标之前对 logits 进行预处理的函数。必须接收 logits 和 labels 作为输入，并返回用于指标计算的 logits。
peft_config (PeftConfig, 可选) — 用于训练时启用 PEFT 的 PEFT 配置。如果为 None，则不使用 PEFT。如果提供，model 将被指定的 PEFT 适配器包装。
formatting_func (Callable，可选) — 用于格式化数据集的函数。必须接收一个示例并返回一个示例。

用于语言模型广义知识蒸馏（GKD）的 Trainer。

有关 GKD 的详细信息，请参阅论文：语言模型的在线策略蒸馏：从自生成错误中学习。

训练

< 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 = None**kwargs )

参数

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

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

GKDConfig

类 trl.experimental.gkd.GKDConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 2e-05lr_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 = 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 = Truelabel_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 = Falserouter_aux_loss_coef: float | None = Nonechat_template_path: str | None = Nonedataset_text_field: str = 'text'dataset_kwargs: dict[str, typing.Any] | str | None = Nonedataset_num_proc: int | None = Noneeos_token: str | None = Nonemax_length: int | None = 1024truncation_mode: str = 'keep_start'shuffle_dataset: bool = Falsepacking: bool = Falsepacking_strategy: str = 'bfd'padding_free: bool = Falsepad_to_multiple_of: int | str | None = 'auto'eval_packing: bool | None = Nonecompletion_only_loss: bool | None = Noneassistant_only_loss: bool = Falseloss_type: str = 'nll'activation_offloading: bool = Falsepad_token: str | None = Nonetemperature: float = 0.9lmbda: float = 0.5beta: float = 0.5max_new_tokens: int = 128teacher_model_name_or_path: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Nonedisable_dropout: bool = Trueseq_kd: bool = False )

参数

temperature (float，可选，默认为 0.9) — 用于采样的温度。温度越高，生成的补全内容越随机。
lmbda (float，可选，默认为 0.5) — Lambda 参数，用于控制学生数据比例（即 on-policy 学生生成输出的占比）。
beta (float, 可选，默认为 0.5) — 广义 Jensen-Shannon 散度损失在 0.0 和 1.0 之间的插值系数。当 beta 为 0.0 时，损失为 KL 散度。当 beta 为 1.0 时，损失为逆 KL 散度。
max_new_tokens (int，可选，默认为 128) — 每次补全生成的最大 token 数。
teacher_model_name_or_path (str, 可选) — 教师模型的名称或路径。如果为 None，则教师模型将与正在训练的模型相同。
teacher_model_init_kwargs (dict[str, Any], 可选) — 从字符串实例化教师模型时，传递给 AutoModelForCausalLM.from_pretrained 的关键字参数
disable_dropout (bool，可选，默认为 True) — 是否禁用模型中的 dropout。
seq_kd (bool，可选，默认为 False) — Seq_kd 参数，用于控制是否执行序列级 KD（可视为在教师模型生成的输出上进行监督式 FT）。

experimental.gkd.GKDTrainer 的配置类。

此类仅包含 GKD 训练特有的参数。完整的训练参数列表，请参阅 TrainingArguments 和 SFTConfig 的文档。

Update on GitHub

TRL

广义知识蒸馏训练器

概述

使用技巧

预期的数据集类型

GKDTrainer

类 trl.experimental.gkd.GKDTrainer

训练

save_model

push_to_hub

GKDConfig

类 trl.experimental.gkd.GKDConfig