Generalized Knowledge Distillation (GKD) was proposed in On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes by Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem.
The abstract from the paper is the following:
Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher’s distribution. Furthermore, GKD facilitates the seamless integration of distillation with RL fine-tuning (RLHF). We demonstrate the efficacy of GKD for distilling auto-regressive language models on summarization, translation, and arithmetic reasoning tasks, and task-agnostic distillation for instruction-tuning.
The key aspects of GKD are:
This post-training method was contributed by Kashif Rasul and Lewis Tunstall.
The experimental.gkd.GKDTrainer is a wrapper around the SFTTrainer class that takes in a teacher model argument. It needs three parameters to be set via the experimental.gkd.GKDConfig namely:
lmbda: controls the student data fraction, i.e., the proportion of on-policy student-generated outputs. When lmbda=0.0, the loss reduces to supervised JSD where the student is trained with the token-level probabilities of the teacher. When lmbda=1.0, the loss reduces to on-policy JSD, where the student generates output sequences and token-specific feedback on these sequences from the teacher. For values in between [0, 1] it is random between the two based on the lmbda value for each batch.seq_kd: controls whether to perform Sequence-Level KD (can be viewed as supervised FT on teacher-generated out). When seq_kd=True and lmbda=0.0, the loss reduces to supervised JSD, where the teacher generates output sequences and the student receives token-specific feedback on these sequences from the teacher.beta: controls the interpolation in the generalized Jensen-Shannon Divergence. When beta=0.0 the loss approximates forward KL divergence, while for beta=1.0 the loss approximates reverse KL divergence. For values in between [0, 1] it interpolates between the two.The authors find that on-policy data (high lmbda) performs better and the optimal beta varied depending on the task and evaluation method.
Make sure that
attn_implementation="kernels-community/flash-attn2"when training Gemma models. Otherwise you will encounter NaNs in the logits due to the soft capping technique adopted by this architecture.
The basic API is as follows:
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()The dataset should be formatted as a list of “messages” where each message is a list of dictionaries with the following keys:
role: either system, assistant or usercontent: the message content( 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 = (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 )
Parameters
torch.nn.Module or str, optional) —
Model to be trained, or the string identifier of the model to be instantiated from a pretrained model.torch.nn.Module or str, optional) —
Teacher model for knowledge distillation, or the string identifier of the model to be instantiated from a
pretrained model.DataCollator, optional) —
Data collator to batch samples from the dataset. It defaults to a
experimental.utils.DataCollatorForChatML using the processing_class.dict of Dataset, optional) —
Dataset for evaluation.Callable, optional) —
Function to compute metrics at evaluation. Must take in an EvalPrediction and return a
dictionary string to float.list of TrainerCallback, optional) —
Callbacks to use during training.tuple of torch.optim.Optimizer and torch.optim.lr_scheduler.LambdaLR, optional, defaults to (None, None)) —
Tuple containing the optimizer and the learning rate scheduler to use for training.Callable, optional) —
Function to preprocess the logits before computing the metrics. Must take in the logits and labels and
return the logits to be used for metrics computation.None, PEFT is not used. If provided, the model will be
wrapped with the specified PEFT adapter.Callable, optional) —
Function to format the dataset. Must take in an example and return an example.Trainer for Generalized Knowledge Distillation (GKD) of language models.
For details on GKD, see the paper: On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.
( resume_from_checkpoint: str | bool | None = Nonetrial: optuna.Trial | dict[str, Any] | None = Noneignore_keys_for_eval: list[str] | None = None ) → ~trainer_utils.TrainOutput
Parameters
str or bool, optional) —
If a str, local path to a saved checkpoint as saved by a previous instance of Trainer. If a
bool and equals True, load the last checkpoint in args.output_dir as saved by a previous instance
of Trainer. If present, training will resume from the model/optimizer/scheduler states loaded here.optuna.Trial or dict[str, Any], optional) —
The trial run or the hyperparameter dictionary for hyperparameter search.list[str], optional) —
A list of keys in the output of your model (if it is a dictionary) that should be ignored when
gathering predictions for evaluation during the training.Returns
~trainer_utils.TrainOutput
Object containing the global step count, training loss, and metrics.
Main training entry point.
Will save the model, so you can reload it using from_pretrained().
Will only save from the main process.
( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )
Parameters
str, optional, defaults to "End of training") —
Message to commit while pushing.bool, optional, defaults to True) —
Whether the function should return only when the git push has finished.str, optional, defaults to None) —
Token with write permission to overwrite Trainer’s original args.str, optional) —
The git revision to commit from. Defaults to the head of the “main” branch.dict[str, Any], optional) —
Additional keyword arguments passed along to ~Trainer.create_model_card.Upload self.model and self.processing_class to the 🤗 model hub on the repo self.args.hub_model_id.
( 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 = <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 | None = Noneeval_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 )
Parameters
float, optional, defaults to 0.9) —
Temperature for sampling. The higher the temperature, the more random the completions.float, optional, defaults to 0.5) —
Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy
student-generated outputs).float, optional, defaults to 0.5) —
Interpolation coefficient between 0.0 and 1.0 of the Generalized Jensen-Shannon Divergence loss. When
beta is 0.0, the loss is the KL divergence. When beta is 1.0, the loss is the Inverse KL Divergence.int, optional, defaults to 128) —
Maximum number of tokens to generate per completion.str, optional) —
Model name or path of the teacher model. If None, the teacher model will be the same as the model being
trained.dict[str, Any], optional) —
Keyword arguments to pass to AutoModelForCausalLM.from_pretrained when instantiating the teacher model
from a string.bool, optional, defaults to True) —
Whether to disable dropout in the model.bool, optional, defaults to False) —
Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on
teacher-generated output).Configuration class for experimental.gkd.GKDTrainer.
This class includes only the parameters that are specific to GKD training. For a full list of training arguments, please refer to the TrainingArguments and SFTConfig documentation.