TRL documentation
IW-OPD
IW-OPD
On the Position Bias of On-Policy Distillation 논문에서 저자들은 Importance-Weighted On-Policy Distillation(IW-OPD)을 소개합니다. IW-OPD는 누적된 teacher-student prefix 불일치에 따라 샘플링된 토큰의 업데이트를 재가중함으로써 on-policy distillation의 위치 편향을 해결합니다. 초기 토큰은 더 큰 가중치를 유지하는 반면, 높은 drift 이후의 후기 토큰은 가중치가 낮아집니다.
IW-OPD를 사용하려면 trl.experimental.iw_opd의 experimental.iw_opd.IWOPDTrainer 클래스를 사용할 수 있습니다.
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는 온-폴리시 목적 함수입니다: distillation_objective="iw_opd"는 lmbda=1.0(기본값)을 요구합니다. iw_opd_gamma는 해당 논문의 알고리즘 1에서 정의된 중요도 가중치 증폭입니다.
IWOPDTrainer
클래스 trl.experimental.iw_opd.IWOPDTrainer
< 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.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 )
중요도 가중치 기반 온-폴리시 증류(Importance-Weighted On-Policy Distillation, IW-OPD)를 위한 Trainer
IW-OPD(On the Position Bias of On-Policy Distillation)는 누적된 teacher-student prefix 불일치에 따라 샘플링된 토큰의 distillation 신호를 재가중합니다. 이를 통해 supervision이 분포에서 벗어난 후속 토큰의 가중치를 낮춥니다. distillation_objective=“iw_opd”를 선택하여 사용하세요.
이 트레이너는 리팩토링 이전의
DistillationTrainer를 동결한 스냅샷으로, IW-OPD(안정적이고 전체 어휘를 사용하는DistillationTrainer에 적합하지 않은 교사 기반 정책 그라디언트 방법)의 저장소로 유지됩니다. 이 트레이너는 능동적으로 유지 관리되지 않으며 해당 트레이너의 후속 개선 사항을 반영하지 않습니다.
지원:
- 일반화된 JSD 손실(beta를 통한 forward KL, reverse KL 또는 보간된 JSD)
- 샘플링된 토큰 IW-OPD 목적 함수(distillation_objective=“iw_opd”)
- lmbda를 통한 온-폴리시/오프-폴리시 혼합 (그라디언트 축적에 걸쳐 버퍼링)
- 로컬 teacher 모델 또는 vLLM 서버를 통한 외부 teacher
- vLLM 또는 model.generate()를 통한 Student의 on-policy 생성
train
< source >( 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와 같으면, 이전Trainer인스턴스가 args.output_dir에 저장한 마지막 체크포인트를 로드합니다. 이 값이 지정되어 있으면, 여기서 로드된 모델/옵티마이저/스케줄러 상태로부터 학습이 재개됩니다. - trial (
optuna.Trial또는dict[str, Any], 선택 사항) — 하이퍼파라미터 탐색을 위한 트라이얼 실행 또는 하이퍼파라미터 사전입니다. - ignore_keys_for_eval (
list[str], 선택 사항) — 모델의 출력(사전에 해당하는 경우)에서 학습 중 평가 예측을 수집할 때 무시해야 할 키의 목록입니다.
반환값
~trainer_utils.TrainOutput
전체 스텝 수, 학습 손실, 메트릭을 포함하는 객체.
주요 학습 진입점.
모델을 저장하므로 from_pretrained()를 사용하여 다시 로드할 수 있습니다.
main 프로세스에서만 저장합니다.
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) — Trainer의 원래 args를 덮어쓸 수 있는 쓰기 권한이 있는 토큰. - revision (
str, 선택 사항) — 커밋할 git revision입니다. 기본값은 “main” 브랜치의 head입니다. - kwargs (
dict[str, Any], 선택) —~Trainer.create_model_card에 함께 전달되는 추가 키워드 인자입니다.
self.args.hub_model_id 저장소에 self.model과 self.processing_class를 🤗 모델 허브에 업로드합니다.
IWOPDConfig
클래스 trl.experimental.iw_opd.IWOPDConfig
< source >( 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], 선택 사항) —model인수가 문자열로 제공될 때AutoModelForCausalLM.from_pretrained에 사용되는 키워드 인수입니다.revision값은 처리 클래스를 로드할 때에도 사용됩니다. - trust_remote_code (
bool, optional, defaults toFalse) — 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) — 각 gradient accumulation slice에서 on-policy(학생이 생성한) 데이터를 사용할 확률입니다.0.0은 완전히 off-policy(데이터셋의 completions만 사용)를 의미하며,1.0은 완전히 on-policy를 의미합니다. - beta (
float, 선택 사항, 기본값1.0) — Generalized Jensen-Shannon Divergence 손실의 보간 계수입니다.0.0이면 손실은 forward KL 발산이 됩니다.1.0이면 손실은 reverse 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") —beta > 0이고loss_top_k == 1인 경우 reverse-KL top-1 토큰을 선택하는 규칙입니다."sampled"는 배치 내 실제 completion 토큰을 사용합니다."argmax"는 student의 확률이 가장 높은 토큰을 사용합니다. 이 설정은 항상 teacher의 top-1 토큰을 사용하는 forward-KL support에는 영향을 주지 않습니다.beta == 0이거나loss_top_k != 1인 경우 무시됩니다. - max_completion_length (
int, 선택 사항, 기본값512) — 온-폴리시(on-policy) 생성 시 완료(completion)당 생성할 최대 토큰 수. - max_prompt_length (
int또는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또는None, 선택 사항) — 교사 모델의 리비전(예: 브랜치 이름, 태그, 커밋 해시). - teacher_model_init_kwargs (
dict[str, Any]또는None, 선택 사항) — 문자열에서 teacher model을 인스턴스화할 때AutoModelForCausalLM.from_pretrained에 전달되는 키워드 인자 - use_teacher_server (
bool, 선택 사항, 기본값False) — 로컬 teacher 모델 대신 외부 vLLM teacher 서버를 사용할지 여부입니다. - teacher_model_server_url (
str또는None, 선택 사항) — teacher 모델을 호스팅하는 vLLM 서버의 Base URL(예:"http://localhost:8000").use_teacher_server=True인 경우, 로컬 forward pass를 실행하는 대신 서버에서 teacher logprobs를 가져옵니다. - loss_top_k (
int, 선택 사항, 기본값1) — JSD/KL 손실을 계산할 때 사용할 상위 토큰 수입니다. 학생과 교사의 분포 모두 이 K개의 토큰으로 제한된 후, 발산을 계산하기 전에 재정규화됩니다. 0인 경우 전체 어휘가 사용됩니다. 로컬 교사의 경우 일반적인 지원 규칙은 forward KL의 경우 교사의 top-k, reverse KL의 경우 학생의 top-k, 혼합 JSD의 경우 두 집합의 합집합입니다.beta > 0이고loss_top_k == 1인 경우, forward 지원은 여전히 교사의 top-1 토큰을 사용하지만, reverse top-1 토큰은reverse_kl_top_1_mode에 의해 제어됩니다.use_teacher_server=True인 경우, 순수 forward 경로(beta=0)는 이 값이 양수여야 하며 forward 항에 교사의 top-k logprobs를 사용합니다.beta > 0인 경우, 서버 기반 증류는loss_top_k == 1을 요구하며"sampled"reverse top-1 토큰만 지원합니다. - loss_add_tail (
bool, 선택 사항, 기본값True) — 손실을 계산할 때, 선택된 top-k 지원 범위 밖의 나머지 확률 질량을 나타내는 tail 버킷을 추가할지 여부입니다.
온-폴리시(on-policy) 생성을 제어하는 매개변수
- num_generations (
int, 선택 사항, 기본값1) — on-policy 생성 과정에서 프롬프트당 생성할 완성본의 수. - generation_batch_size (
int또는None, 선택 사항) — 옵티마이저 스텝당 워커당 고유 프롬프트 수입니다.None인 경우(per_device_train_batch_size * gradient_accumulation_steps) // num_generations로 계산됩니다. - top_p (
float, 선택 사항, 기본값0.95) — 온-폴리시(on-policy) 생성을 위한 Top-p (nucleus) 샘플링 파라미터. - top_k (
int, 선택 사항, 기본값0) — 온-폴리시(on-policy) 생성을 위한 Top-k 샘플링 매개변수입니다.0은 top-k 필터링을 비활성화합니다.
학생 생성을 위한 vLLM을 제어하는 매개변수
- use_vllm (
bool, 선택 사항, 기본값False) — 학생 모델에서 온-폴리시(on-policy) 완성물을 생성할 때 vLLM을 사용할지 여부입니다. - vllm_mode (
str, 선택 사항, 기본값"colocate") — student vLLM 통합을 위한 모드."server"또는"colocate"중 하나입니다. - vllm_server_base_url (
str또는None, 선택 사항) — student vLLM 서버의 기본 URL입니다. 이 값이 제공되면vllm_server_host와vllm_server_port는 무시됩니다. - vllm_server_host (
str, 선택 사항, 기본값"0.0.0.0") — 학생 vLLM 서버의 호스트입니다. - vllm_server_port (
int, 선택 사항, 기본값8001) — student vLLM 서버의 포트입니다. - vllm_server_timeout (
float, 선택 사항, 기본값240.0) — student vLLM 서버에 연결하는 데 걸리는 시간 제한입니다. - vllm_group_port (
int, 선택 사항, 기본값51216) — vLLM 가중치 업데이트 그룹(NCCL communicator)에 사용할 포트입니다. - vllm_gpu_memory_utilization (
float, 선택 사항, 기본값0.3) — 동일 위치에 배치된 student vLLM 엔진의 GPU 메모리 사용률입니다. - vllm_tensor_parallel_size (
int, 선택 사항, 기본값1) — 동일 위치에 배치된 student vLLM 엔진의 Tensor parallel size입니다. - vllm_max_model_length (
int또는None, 선택 사항) — 동일 위치에 배치된 vLLM 엔진의 최대 모델 시퀀스 길이입니다. - vllm_model_impl (
str, 선택 사항, 기본값"vllm") — vLLM의 모델 구현 백엔드입니다."vllm"또는"transformers"를 사용하세요. - vllm_structured_outputs_regex (
str또는None, 선택 사항) — vLLM structured outputs에 사용할 정규식 패턴입니다. - vllm_sync_frequency (
int, 선택 사항, 기본값1) — vLLM 엔진에 학생 모델 가중치를 동기화할 주기(학습 스텝 단위). - vllm_enable_sleep_mode (
bool, 선택 사항, 기본값False) — 옵티마이저 단계에서 학생 가중치를 오프로드하기 위해 vLLM sleep mode를 활성화합니다.
로깅을 제어하는 매개변수
- log_completions (
bool, 선택 사항, 기본값False) —log_completions_steps스텝마다 (prompt, completion) 쌍의 샘플을 기록할지 여부입니다.rich가 설치되어 있으면 샘플을 출력합니다.wandb및/또는trackio로깅이 활성화되어 있으면wandb및/또는trackio에 기록합니다. - log_completions_steps (
int, 선택 사항, 기본값100) — completions 로깅 사이의 스텝 수.log_completions가True인 경우에만 사용됩니다. - num_completions_to_print (
int또는None, 선택 사항) — 출력할 completions의 수입니다.None인 경우 모든 completions이 기록됩니다.
experimental.iw_opd.IWOPDTrainer의 구성 클래스입니다.
지식 증류에 특화된 매개변수로 TrainingArguments를 확장합니다. 이 설정은 SFTConfig와 독립적이며, 필요한 모든 필드가 여기에 선언되어 있습니다.
HfArgumentParser를 사용하면 이 클래스를 커맨드 라인에서 지정할 수 있는 argparse 인수로 변환할 수 있습니다.