Hugging Face's logo

TRL documentation

IW-OPD

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→