Hugging Face's logo

TRL documentation

IW-OPD

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.13.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

IW-OPD

In the paper On the Position Bias of On-Policy Distillation, the authors introduce Importance-Weighted On-Policy Distillation (IW-OPD). IW-OPD addresses position bias in on-policy distillation by reweighting sampled-token updates according to accumulated teacher-student prefix discrepancy. Early tokens keep larger weights, while later tokens after high drift are downweighted.

To use IW-OPD, you can use the experimental.iw_opd.IWOPDTrainer class in trl.experimental.iw_opd.

IW-OPD is currently part of the trl.experimental namespace. APIs may change without notice while the feature is iterated on.

Usage

from trl.experimental.iw_opd import IWOPDConfig, IWOPDTrainer

training_args = IWOPDConfig(
    distillation_objective="iw_opd",
    iw_opd_gamma=0.5,
)
trainer = IWOPDTrainer(
    model="Qwen/Qwen3-0.6B",
    teacher_model="...",
    train_dataset=...,
    args=training_args,
)
trainer.train()

IW-OPD is an on-policy objective: distillation_objective="iw_opd" requires lmbda=1.0 (the default). iw_opd_gamma is the importance-weight amplification from Algorithm 1 of the paper.

IWOPDTrainer

class trl.experimental.iw_opd.IWOPDTrainer

< >

( 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.iw_opd.iw_opd_config.IWOPDConfig | 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: typing.Optional[ForwardRef('PeftConfig')] = None )

Trainer for Importance-Weighted On-Policy Distillation (IW-OPD).

IW-OPD (On the Position Bias of On-Policy Distillation) reweights the sampled-token distillation signal by accumulated teacher-student prefix discrepancy, downweighting later tokens whose supervision has drifted out of distribution. Select it with distillation_objective=“iw_opd”.

This trainer is a frozen snapshot of the pre-refactor DistillationTrainer, kept as the home for IW-OPD (a teacher-guided policy-gradient method that does not fit the stable, full-vocabulary DistillationTrainer). It is not actively maintained and will not track that trainer’s later improvements.

Supports:

  • Generalized JSD loss (forward KL, reverse KL, or interpolated JSD via beta)
  • Sampled-token IW-OPD objective (distillation_objective=“iw_opd”)
  • On-policy / off-policy mixing via lmbda (buffered across gradient accumulation)
  • Local teacher model or external teacher via vLLM server
  • Student on-policy generation via vLLM or model.generate()
  • Liger kernel for memory-efficient fused JSD loss

train

< >

( 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

  • resume_from_checkpoint (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.
  • trial (optuna.Trial or dict[str, Any], optional) — The trial run or the hyperparameter dictionary for hyperparameter search.
  • ignore_keys_for_eval (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.

save_model

< >

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

Will save the model, so you can reload it using from_pretrained().

Will only save from the main process.

push_to_hub

< >

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

Parameters

  • commit_message (str, optional, defaults to "End of training") — Message to commit while pushing.
  • blocking (bool, optional, defaults to True) — Whether the function should return only when the git push has finished.
  • token (str, optional, defaults to None) — Token with write permission to overwrite Trainer’s original args.
  • revision (str, optional) — The git revision to commit from. Defaults to the head of the “main” branch.
  • kwargs (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.

IWOPDConfig

class trl.experimental.iw_opd.IWOPDConfig

< >

( 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 = 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 = Falsemax_length: int | None = 1024temperature: float = 1.0lmbda: float = 1.0beta: float = 1.0distillation_objective: str = 'iw_opd'iw_opd_gamma: float = 0.5iw_opd_epsilon: float = 1e-08reverse_kl_top_1_mode: str = 'sampled'max_completion_length: int = 512max_prompt_length: int | None = Nonedisable_dropout: bool = Trueteacher_model_name_or_path: str | None = Noneteacher_model_revision: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Noneuse_teacher_server: bool = Falseteacher_model_server_url: str | None = Noneloss_top_k: int = 1loss_add_tail: bool = Truenum_generations: int = 1generation_batch_size: int | None = Nonetop_p: float = 0.95top_k: int = 0use_vllm: bool = Falsevllm_mode: str = 'colocate'vllm_server_base_url: str | None = Nonevllm_server_host: str = '0.0.0.0'vllm_server_port: int = 8001vllm_server_timeout: float = 240.0vllm_group_port: int = 51216vllm_gpu_memory_utilization: float = 0.3vllm_tensor_parallel_size: int = 1vllm_max_model_length: int | None = Nonevllm_model_impl: str = 'vllm'vllm_structured_outputs_regex: str | None = Nonevllm_sync_frequency: int = 1vllm_enable_sleep_mode: bool = Falselog_completions: bool = Falselog_completions_steps: int = 100num_completions_to_print: int | None = None )

Parameters that control the model

  • model_init_kwargs (dict[str, Any], optional) — Keyword arguments for AutoModelForCausalLM.from_pretrained, used when the model argument of the trainer is provided as a string. The revision value is also used when loading the processing class.
  • 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.
  • max_length (int or None, optional, defaults to 1024) — Maximum total sequence length (prompt + completion) for tokenization and truncation.

Parameters that control the distillation

  • temperature (float, optional, defaults to 1.0) — Temperature for sampling during generation and for computing the distillation loss. Higher values produce softer probability distributions.
  • lmbda (float, optional, defaults to 1.0) — Probability of using on-policy (student-generated) data for each gradient accumulation slice. A value of 0.0 means fully off-policy (dataset completions only), 1.0 means fully on-policy.
  • beta (float, optional, defaults to 1.0) — Interpolation coefficient for the Generalized Jensen-Shannon Divergence loss. When 0.0, the loss is the forward KL divergence. When 1.0, the loss is the reverse KL divergence. When 0.5, it is the standard JSD.
  • distillation_objective (str, optional, defaults to "iw_opd") — Objective to optimize. "iw_opd" uses the sampled-token Importance-Weighted On-Policy Distillation objective. "jsd" keeps the generalized JSD/KL objective.
  • iw_opd_gamma (float, optional, defaults to 0.5) — Importance-weight amplification for distillation_objective="iw_opd".
  • iw_opd_epsilon (float, optional, defaults to 1e-8) — Stabilizer used when normalizing IW-OPD prefix weights.
  • reverse_kl_top_1_mode (str, optional, defaults to "sampled") — Selection rule for the reverse-KL top-1 token when beta > 0 and loss_top_k == 1. "sampled" uses the actual completion token in the batch. "argmax" uses the student’s highest-probability token. This setting does not affect the forward-KL support, which always uses the teacher’s top-1 token. Ignored when beta == 0 or loss_top_k != 1.
  • max_completion_length (int, optional, defaults to 512) — Maximum number of tokens to generate per completion during on-policy generation.
  • max_prompt_length (int or None, optional) — Maximum number of tokens for the prompt. If None, auto-computed as max_length - max_completion_length. Prompts are truncated according to the tokenizer’s truncation_side setting.
  • disable_dropout (bool, optional, defaults to True) — Whether to disable dropout in the student model during training.

Parameters that control the teacher model

  • teacher_model_name_or_path (str or None, optional) — Model name or path for the teacher model. Used when the teacher is loaded locally.
  • teacher_model_revision (str or None, optional) — Model revision of the teacher model (e.g., branch name, tag, or commit hash).
  • teacher_model_init_kwargs (dict[str, Any] or None, optional) — Keyword arguments passed to AutoModelForCausalLM.from_pretrained when instantiating the teacher model from a string.
  • use_teacher_server (bool, optional, defaults to False) — Whether to use an external vLLM teacher server instead of a local teacher model.
  • teacher_model_server_url (str or None, optional) — Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). When set, teacher logprobs are fetched from the server instead of running a local forward pass when use_teacher_server=True.
  • loss_top_k (int, optional, defaults to 1) — Number of top tokens to use when computing the JSD/KL loss. Both student and teacher distributions are restricted to these K tokens and re-normalized before computing divergence. If 0, the full vocabulary is used. For local teachers, the general support rule is teacher top-k for forward KL, student top-k for reverse KL, and the union for mixed JSD. When beta > 0 and loss_top_k == 1, the forward support still uses the teacher’s top-1 token, while the reverse top-1 token is controlled by reverse_kl_top_1_mode. When use_teacher_server=True, the pure forward path (beta=0) requires this to be positive and uses the teacher’s top-k logprobs for the forward term. When beta > 0, server-backed distillation requires loss_top_k == 1 and only supports "sampled" reverse top-1 tokens.
  • loss_add_tail (bool, optional, defaults to True) — Whether to append a tail bucket that represents the remaining probability mass outside the selected top-k support when computing the loss.

Parameters that control on-policy generation

  • num_generations (int, optional, defaults to 1) — Number of completions to generate per prompt during on-policy generation.
  • generation_batch_size (int or None, optional) — Number of unique prompts per worker per optimizer step. If None, computed from (per_device_train_batch_size * gradient_accumulation_steps) // num_generations.
  • top_p (float, optional, defaults to 0.95) — Top-p (nucleus) sampling parameter for on-policy generation.
  • top_k (int, optional, defaults to 0) — Top-k sampling parameter for on-policy generation. 0 disables top-k filtering.

Parameters that control vLLM for student generation

  • use_vllm (bool, optional, defaults to False) — Whether to use vLLM for generating on-policy completions from the student model.
  • vllm_mode (str, optional, defaults to "colocate") — Mode for student vLLM integration. Either "server" or "colocate".
  • vllm_server_base_url (str or None, optional) — Base URL for the student vLLM server. If provided, vllm_server_host and vllm_server_port are ignored.
  • vllm_server_host (str, optional, defaults to "0.0.0.0") — Host of the student vLLM server.
  • vllm_server_port (int, optional, defaults to 8001) — Port of the student vLLM server.
  • vllm_server_timeout (float, optional, defaults to 240.0) — Timeout for connecting to the student vLLM server.
  • vllm_group_port (int, optional, defaults to 51216) — Port for the vLLM weight-update group (NCCL communicator).
  • vllm_gpu_memory_utilization (float, optional, defaults to 0.3) — GPU memory utilization for the colocated student vLLM engine.
  • vllm_tensor_parallel_size (int, optional, defaults to 1) — Tensor parallel size for the colocated student vLLM engine.
  • vllm_max_model_length (int or None, optional) — Maximum model sequence length for the colocated vLLM engine.
  • vllm_model_impl (str, optional, defaults to "vllm") — Model implementation backend for vLLM. Use "vllm" or "transformers".
  • vllm_structured_outputs_regex (str or None, optional) — Regex pattern for vLLM structured outputs.
  • vllm_sync_frequency (int, optional, defaults to 1) — Frequency (in training steps) to synchronize student model weights to the vLLM engine.
  • vllm_enable_sleep_mode (bool, optional, defaults to False) — Enable vLLM sleep mode to offload student weights during the optimizer step.

Parameters that control logging

  • log_completions (bool, optional, defaults to False) — Whether to log a sample of (prompt, completion) pairs every log_completions_steps steps. If rich is installed, it prints the sample. If wandb and/or trackio logging is enabled, it logs it to wandb and/or trackio.
  • log_completions_steps (int, optional, defaults to 100) — Number of steps between logging completions. Only used if log_completions is True.
  • num_completions_to_print (int or None, optional) — Number of completions to print. If None, all completions are logged.

Configuration class for the experimental.iw_opd.IWOPDTrainer.

Extends TrainingArguments with parameters specific to knowledge distillation. This config is independent of SFTConfig — all necessary fields are declared here.

Using HfArgumentParser we can turn this class into argparse arguments that can be specified on the command line.

Update on GitHub

←GSPO-token MergeModelCallback→