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.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

IW-OPD

В статье On the Position Bias of On-Policy Distillation авторы представляют Importance-Weighted On-Policy Distillation (IW-OPD). IW-OPD устраняет позиционное смещение в on-policy дистилляции, пересчитывая веса обновлений для сэмплированных токенов на основе накопленного расхождения префиксов между учителем и студентом. Ранние токены сохраняют более высокие веса, тогда как поздние токены после значительного дрейфа получают пониженные веса.

Для использования IW-OPD можно применить класс experimental.iw_opd.IWOPDTrainer из модуля trl.experimental.iw_opd.

IW-OPD в настоящее время входит в пространство имён trl.experimental. API могут измениться без предварительного уведомления, пока над функцией ведётся работа.

Использование

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 — это on-policy целевая функция: distillation_objective="iw_opd" требует lmbda=1.0 (значение по умолчанию). iw_opd_gamma — это усиление важности (importance-weight amplification) из Алгоритма 1 статьи.

IWOPDTrainer

класс 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[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: typing.Optional[ForwardRef('PeftConfig')] = None )

Trainer для дистилляции с важностным взвешиванием по текущей политике (IW-OPD).

IW-OPD (On the Position Bias of On-Policy Distillation) пересчитывает сигнал дистилляции для сэмплированных токенов, используя накопленное расхождение префиксов учителя и ученика, и снижает вес поздних токенов, чей сигнал обучения ушел за пределы распределения. Выберите его с помощью distillation_objective=“iw_opd”.

Этот тренажёр представляет собой замороженную копию DistillationTrainer до рефакторинга и служит местом для IW-OPD (метода политического градиента с наставником, который не вписывается в стабильный DistillationTrainer с полным словарём). Он не поддерживается активно и не будет учитывать последующие улучшения того тренажёра.

Поддерживает:

  • Обобщённая функция потерь JSD (прямая KL, обратная KL или интерполированная JSD через beta)
  • Целевая функция IW-OPD для сэмплированных токенов (distillation_objective=“iw_opd”)
  • Смешивание on-policy / off-policy через lmbda (буферизуется при накоплении градиентов)
  • Локальная учительская модель или внешняя учительская модель через сервер vLLM
  • Генерация студента по его собственной политике с помощью vLLM или model.generate()

обучение

< >

( resume_from_checkpoint: str | bool | None = Nonetrial: optuna.Trial | dict[str, Any] | None = Noneignore_keys_for_eval: list[str] | None = None ) → ~trainer_utils.TrainOutput

Параметры

  • resume_from_checkpoint (str или bool, необязательный) — Если передано значение типа str, это локальный путь к сохранённой контрольной точке, созданной предыдущим экземпляром Trainer. Если передано значение типа bool и оно равно True, загружается последняя контрольная точка из args.output_dir, сохранённая предыдущим экземпляром Trainer. Если параметр указан, обучение возобновляется из состояний модели/оптимизатора/планировщика, загруженных из этой контрольной точки.
  • trial (optuna.Trial или dict[str, Any], необязательный) — Запуск испытания или словарь гиперпараметров для поиска гиперпараметров.
  • ignore_keys_for_eval (list[str], необязательный) — Список ключей в выводе вашей модели (если он представляет собой словарь), которые следует игнорировать при сборе предсказаний для оценки во время обучения.

Возвращает

~trainer_utils.TrainOutput

Объект, содержащий количество глобальных шагов, потерю обучения и метрики.

Основная точка входа для обучения.

save_model

< >

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

Сохранит модель, чтобы вы могли загрузить её снова с помощью from_pretrained().

Сохранение будет выполняться только из основного процесса.

push_to_hub

< >

( 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) — Токен с правом записи для перезаписи исходных аргументов Trainer.
  • revision (str, необязательный) — Git-ревизия, от которой выполняется коммит. По умолчанию используется head ветки “main”.
  • kwargs (dict[str, Any], необязательный) — Дополнительные именованные аргументы, передаваемые в ~Trainer.create_model_card.

Загрузите self.model и self.processing_class в модельный хаб 🤗 в репозиторий self.args.hub_model_id.

IWOPDConfig

класс 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[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 = 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 )

Параметры, управляющие моделью

  • model_init_kwargs (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.
  • max_length (int или None, необязательный, по умолчанию 1024) — Максимальная общая длина последовательности (запрос + завершение) для токенизации и усечения.

Параметры, управляющие дистилляцией

  • temperature (float, необязательный, по умолчанию 1.0) — Температура для выборки при генерации и для вычисления потерь дистилляции. Более высокие значения приводят к более мягким распределениям вероятностей.
  • lmbda (float, необязательный, по умолчанию 1.0) — Вероятность использования on-policy (сгенерированных студентом) данных для каждого среза накопления градиентов. Значение 0.0 означает полностью off-policy (только завершения из набора данных), 1.0 — полностью on-policy.
  • beta (float, необязательный, по умолчанию 1.0) — Коэффициент интерполяции для функции потерь обобщённого расхождения Дженсена–Шаннона. При 0.0 функция потерь представляет собой прямое расхождение KL. При 1.0 функция потерь представляет собой обратное расхождение KL. При 0.5 это стандартное расхождение JSD.
  • distillation_objective (str, необязательный, по умолчанию "iw_opd") — Целевая функция для оптимизации. "iw_opd" использует целевую функцию Importance-Weighted On-Policy Distillation для сэмплированных токенов. "jsd" сохраняет обобщённую целевую функцию JSD/KL.
  • iw_opd_gamma (float, необязательный, по умолчанию 0.5) — Усиление с весами важности для distillation_objective="iw_opd".
  • iw_opd_epsilon (float, необязательный, по умолчанию 1e-8) — Стабилизатор, используемый при нормализации весов префиксов IW-OPD.
  • reverse_kl_top_1_mode (str, необязательный, по умолчанию "sampled") — Правило выбора токена top-1 для reverse-KL, если beta > 0 и loss_top_k == 1. "sampled" использует фактический токен завершения в батче. "argmax" использует токен с наибольшей вероятностью у ученика. Этот параметр не влияет на поддержку forward-KL, которая всегда использует токен top-1 учителя. Игнорируется, если beta == 0 или loss_top_k != 1.
  • max_completion_length (int, необязательный, по умолчанию 512) — Максимальное количество токенов, генерируемых на каждое завершение во время on-policy генерации.
  • max_prompt_length (int or None, необязательный) — Максимальное количество токенов в промпте. Если None, значение вычисляется автоматически как max_length - max_completion_length. Промпты обрезаются в соответствии с настройкой truncation_side токенизатора.
  • disable_dropout (bool, необязательный, по умолчанию True) — Определяет, следует ли отключать dropout в модели-ученике во время обучения.

Параметры, управляющие моделью-учителем

  • teacher_model_name_or_path (str или None, необязательный) — Имя модели или путь к модели учителя. Используется, когда учитель загружается локально.
  • teacher_model_revision (str or None, необязательный) — Ревизия модели-учителя (например, имя ветки, тег или хеш коммита).
  • teacher_model_init_kwargs (dict[str, Any] или None, необязательный) — Параметры, передаваемые в AutoModelForCausalLM.from_pretrained при создании модели-учителя из строки.
  • use_teacher_server (bool, необязательный, по умолчанию False) — Использовать ли внешний сервер vLLM teacher вместо локальной teacher-модели.
  • teacher_model_server_url (str or None, необязательный) — Базовый URL vLLM-сервера, на котором размещена teacher-модель (например, "http://localhost:8000"). Если задан, logprobs teacher-модели запрашиваются с сервера, а не вычисляются локальным forward-проходом, когда use_teacher_server=True.
  • loss_top_k (int, необязательный, по умолчанию 1) — Количество верхних токенов, используемых при вычислении потерь JSD/KL. Распределения как ученика, так и учителя ограничиваются этими K токенами и нормализуются заново перед вычислением расхождения. Если значение равно 0, используется весь словарь. Для локальных учителей общее правило поддержки: top-k учителя для прямого KL, top-k ученика для обратного KL и объединение для смешанного JSD. Если beta > 0 и loss_top_k == 1, прямая поддержка всё равно использует top-1 токен учителя, тогда как обратный top-1 токен контролируется параметром reverse_kl_top_1_mode. Если use_teacher_server=True, чистый прямой путь (beta=0) требует, чтобы это значение было положительным, и использует top-k logprobs учителя для прямого слагаемого. Если beta > 0, дистилляция с поддержкой сервера требует loss_top_k == 1 и поддерживает только "sampled" для обратных top-1 токенов.
  • loss_add_tail (bool, необязательный, по умолчанию True) — Определяет, следует ли добавлять хвостовое значение, представляющее оставшуюся вероятность вне выбранной top-k области определения, при вычислении функции потерь.

Параметры, управляющие генерацией on-policy

  • num_generations (int, необязательный, по умолчанию 1) — Количество завершений, генерируемых для каждого промпта при on-policy генерации.
  • generation_batch_size (int or None, необязательный) — Количество уникальных промптов на каждый воркер за шаг оптимизатора. Если None, вычисляется как (per_device_train_batch_size * gradient_accumulation_steps) // num_generations.
  • top_p (float, необязательный, по умолчанию 0.95) — Параметр top-p (nucleus) выборки для генерации on-policy.
  • top_k (int, необязательный, по умолчанию 0) — Параметр top-k выборки для генерации on-policy. 0 отключает top-k фильтрацию.

Параметры, управляющие vLLM для генерации студентом

  • use_vllm (bool, необязательный, по умолчанию False) — Использовать ли vLLM для генерации on-policy завершений от модели-ученика.
  • vllm_mode (str, необязательный, по умолчанию "colocate") — Режим интеграции vLLM для student. Допустимые значения: "server" или "colocate".
  • vllm_server_base_url (str или None, необязательный) — Базовый URL для сервера vLLM студента. Если указан, vllm_server_host и vllm_server_port игнорируются.
  • vllm_server_host (str, необязательный, по умолчанию "0.0.0.0") — Хост сервера vLLM студента.
  • vllm_server_port (int, необязательный, по умолчанию 8001) — Порт сервера vLLM студента.
  • vllm_server_timeout (float, необязательный, по умолчанию 240.0) — Время ожидания подключения к серверу vLLM студента.
  • vllm_group_port (int, необязательный, по умолчанию 51216) — Порт для группы обновления весов vLLM (коммутатор NCCL).
  • vllm_gpu_memory_utilization (float, необязательный, по умолчанию 0.3) — доля использования памяти GPU для движка vLLM студента, работающего на том же устройстве.
  • vllm_tensor_parallel_size (int, необязательный, по умолчанию 1) — Размер тензорного параллелизма для локального движка vLLM студента.
  • vllm_max_model_length (int или None, необязательный) — Максимальная длина последовательности модели для совмещённого движка vLLM.
  • vllm_model_impl (str, необязательный, по умолчанию "vllm") — Бэкенд реализации модели для vLLM. Используйте "vllm" или "transformers".
  • vllm_structured_outputs_regex (str или None, необязательный) — Шаблон регулярного выражения для структурированных выходов vLLM.
  • vllm_sync_frequency (int, необязательный, по умолчанию 1) — Частота (в шагах обучения) синхронизации весов модели-ученика с движком vLLM.
  • vllm_enable_sleep_mode (bool, необязательный, по умолчанию False) — Включает режим сна vLLM для выгрузки весов ученика во время шага оптимизатора.

Параметры, управляющие ведением журнала

  • log_completions (bool, необязательный, по умолчанию False) — Определяет, следует ли логировать выборку пар (prompt, completion) каждые log_completions_steps шагов. Если установлен пакет rich, выборка выводится в консоль. Если включено логирование в wandb и/или trackio, данные записываются в wandb и/или trackio.
  • log_completions_steps (int, необязательный, по умолчанию 100) — Количество шагов между записью в журнал результатов генерации. Используется только если log_completions равно True.
  • num_completions_to_print (int or None, необязательный) — Количество завершений для вывода. Если None, в журнал записываются все завершения.

Класс конфигурации для experimental.iw_opd.IWOPDTrainer.

Расширяет TrainingArguments параметрами, специфичными для дистилляции знаний. Эта конфигурация не зависит от SFTConfig — все необходимые поля объявлены здесь.

С помощью HfArgumentParser мы можем превратить этот класс в аргументы argparse, которые можно указать в командной строке.

Update on GitHub

←GOLD MergeModelCallback→