일반화 지식 증류 트레이너

개요

일반화 지식 증류(GKD)는 Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem이 On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes에서 제안했습니다.

해당 논문의 초록은 다음과 같습니다:

Knowledge distillation(KD)은 더 작은 student model을 학습시켜 teacher model의 추론 비용과 메모리 사용량을 줄이기 위해 teacher model을 압축하는 데 널리 사용됩니다. 그러나 현재 auto-regressive sequence model을 위한 KD 방법은 학습 중에 접하는 출력 시퀀스와 student가 추론 중에 생성하는 출력 시퀀스 사이의 분포 불일치 문제를 겪고 있습니다. 이 문제를 해결하기 위해 Generalized Knowledge Distillation(GKD)을 소개합니다. GKD는 고정된 출력 시퀀스 세트에만 의존하는 대신, teacher가 해당 시퀀스에 대해 제공하는 피드백을 활용하여 student가 스스로 생성한 출력 시퀀스로 student를 학습시킵니다. supervised KD 접근법과 달리, GKD는 student와 teacher 사이에 대안적인 loss function을 사용할 수 있는 유연성을 제공하며, 이는 teacher의 분포를 모방할 표현력이 부족한 student의 경우 유용할 수 있습니다. 또한 GKD는 distillation과 RL fine-tuning(RLHF)의 원활한 통합을 촉진합니다. 요약, 번역, 산술 추론 작업 및 instruction-tuning을 위한 task-agnostic distillation에서 auto-regressive language model을 distill하는 GKD의 효능을 입증합니다.

GKD의 주요 측면은 다음과 같습니다:

자기 회귀 시퀀스 모델에서 학습-추론 분포 불일치를 해결하기 위해, 학생 모델을 자체 생성한 출력 시퀀스로 학습시킵니다.
GKD는 학생 모델이 교사를 완전히 모방할 수 있는 용량이 없을 때 유용할 수 있는 일반화 Jensen-Shannon Divergence(JSD)를 통해 학생 모델과 교사 모델 간의 다양한 발산 측정 방식을 선택할 수 있는 유연성을 제공합니다.

이 포스트 트레이닝 방법은 Kashif Rasul와 Lewis Tunstall이 기여했습니다.

사용 팁

experimental.gkd.GKDTrainer는 teacher model 인수를 받는 SFTTrainer 클래스를 감싸는 래퍼입니다. experimental.gkd.GKDConfig를 통해 다음 세 가지 파라미터를 설정해야 합니다:

lmbda: 학생 데이터 비율, 즉 온-폴리시(on-policy) 학생 생성 출력의 비율을 제어합니다. lmbda=0.0인 경우 손실은 지도 JSD로 축소되며, 학생은 교사의 토큰 수준 확률로 학습됩니다. lmbda=1.0인 경우 손실은 온-폴리시 JSD로 축소되며, 학생은 출력 시퀀스를 생성하고 교사가 해당 시퀀스에 대한 토큰별 피드백을 제공합니다. [0, 1] 사이의 값인 경우 각 배치에 대해 lmbda 값에 따라 두 방식 사이에서 무작위로 선택됩니다.
seq_kd: 시퀀스 레벨 KD(Sequence-Level KD)를 수행할지 여부를 제어합니다(교사가 생성한 출력에 대한 감독된 FT로 볼 수 있음). seq_kd=True이고 lmbda=0.0인 경우, 손실은 감독된 JSD로 축소되며, 이때 교사가 출력 시퀀스를 생성하고 학생은 교사로부터 해당 시퀀스에 대한 토큰별 피드백을 받습니다.
beta: 일반화된 Jensen-Shannon Divergence의 보간을 제어합니다. beta=0.0일 때 손실은 forward KL divergence를 근사하고, beta=1.0일 때 손실은 reverse KL divergence를 근사합니다. [0, 1] 사이의 값에서는 두 분산 사이에 보간합니다.

저자들은 온-폴리시 데이터(높은 lmbda)가 더 나은 성능을 보이며, 최적의 beta는 작업과 평가 방법에 따라 달라진다는 것을 발견했습니다.

Gemma 모델을 학습할 때 attn_implementation="kernels-community/flash-attn2"가 설정되어 있는지 확인하세요. 그렇지 않으면 이 아키텍처가 채택한 soft capping 기법으로 인해 logits에서 NaN이 발생하게 됩니다.

기본 API는 다음과 같습니다:

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()

기대되는 데이터셋 유형

데이터셋은 각 메시지가 다음 키를 가진 사전의 목록인 "messages"의 목록으로 포맷되어야 합니다:

role: system, assistant 또는 user
content: 메시지 내용

GKDTrainer

클래스 trl.experimental.gkd.GKDTrainer

< 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.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[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: PeftConfig | None = Noneformatting_func: collections.abc.Callable | None = None )

매개변수

model (PreTrainedModel 또는 torch.nn.Module 또는 str, 선택 사항) — 학습할 모델 또는 사전 학습된 모델에서 인스턴스화할 모델의 문자열 식별자.
teacher_model (PreTrainedModel 또는 torch.nn.Module 또는 str, 선택 사항) — 지식 증류용 teacher model 또는 사전 학습된 모델에서 인스턴스화할 모델의 문자열 식별자.
args (experimental.gkd.GKDConfig, 선택 사항) — 학습 인자입니다.
data_collator (DataCollator, 선택 사항) — 데이터셋에서 샘플을 배치하기 위한 Data collator입니다. 기본적으로 processing_class를 사용하는 experimental.utils.DataCollatorForChatML로 설정됩니다.
train_dataset (Dataset, 선택 사항) — 학습에 사용할 Dataset입니다.
eval_dataset (Dataset 또는 Dataset의 dict, 선택 사항) — 평가에 사용할 Dataset입니다.
processing_class (PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin 또는 ProcessorMixin, 선택 사항) — 데이터를 처리하는 클래스.
compute_metrics (Callable, 선택 사항) — 평가 시 메트릭을 계산하는 함수. EvalPrediction를 입력으로 받아 문자열에서 float으로 매핑된 딕셔너리를 반환해야 합니다.
callbacks (list of TrainerCallback, optional) — 학습 중에 사용할 콜백입니다.
optimizers (tuple of torch.optim.Optimizer and torch.optim.lr_scheduler.LambdaLR, 선택 사항, 기본값 (None, None)) — 학습에 사용할 옵티마이저와 학습률 스케줄러를 포함하는 튜플입니다.
preprocess_logits_for_metrics (Callable, 선택 사항) — 메트릭을 계산하기 전에 logits를 전처리하는 함수. logits와 labels를 입력으로 받아 메트릭 계산에 사용할 logits를 반환해야 합니다.
peft_config (PeftConfig, 선택 사항) — PEFT를 사용하여 학습할 때 사용할 PEFT 구성입니다. None인 경우 PEFT를 사용하지 않습니다. 값이 제공되면 model은 지정된 PEFT 어댑터로 래핑됩니다.
formatting_func (Callable, 선택 사항) — 데이터셋을 포맷팅하는 함수. 예시를 입력받아 예시를 반환해야 합니다.

언어 모델의 일반화 지식 증류(GKD)를 위한 Trainer입니다.

GKD에 대한 자세한 내용은 다음 논문을 참고하세요: 언어 모델의 온-폴리시 증류: 스스로 생성한 실수로부터 배우기.

train

< source >

매개변수

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

전체 스텝 수, 학습 손실, 메트릭을 포함하는 객체.

주요 학습 진입점.

save_model

< source >

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

모델을 저장하므로 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를 🤗 모델 허브에 업로드합니다.

GKDConfig

클래스 trl.experimental.gkd.GKDConfig

< source >

( 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[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 = 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 | str | None = 'auto'eval_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 )

매개변수

temperature (float, 선택 사항, 기본값 0.9) — 샘플링에 사용할 온도입니다. 온도가 높을수록 완성 결과가 더 무작위적으로 생성됩니다.
lmbda (float, 선택 사항, 기본값 0.5) — 학생 데이터 비율(즉, 온-폴리시 학생 생성 출력의 비율)을 제어하는 람다 파라미터입니다.
beta (float, 선택 사항, 기본값 0.5) — Generalized Jensen-Shannon Divergence 손실의 0.0과 1.0 사이의 보간 계수입니다. beta가 0.0이면 손실은 KL 발산이 됩니다. beta가 1.0이면 손실은 Inverse KL Divergence가 됩니다.
max_new_tokens (int, 선택 사항, 기본값 128) — 완성당 생성할 최대 토큰 수.
teacher_model_name_or_path (str, 선택 사항) — 교사 모델의 이름 또는 경로입니다. None인 경우, 교사 모델은 학습 중인 모델과 동일합니다.
teacher_model_init_kwargs (dict[str, Any], 선택 사항) — 문자열에서 teacher model을 인스턴스화할 때 AutoModelForCausalLM.from_pretrained에 전달할 키워드 인자
disable_dropout (bool, 선택 사항, 기본값 True) — 모델에서 dropout을 비활성화할지 여부입니다.
seq_kd (bool, 선택 사항, 기본값 False) — Sequence-Level KD 수행 여부를 제어하는 Seq_kd 매개변수입니다(교사가 생성한 출력에 대한 supervised FT로 볼 수 있습니다).

experimental.gkd.GKDTrainer의 구성 클래스입니다.

이 클래스에는 GKD 학습에 특화된 매개변수만 포함되어 있습니다. 전체 학습 인자 목록은 TrainingArguments 및 SFTConfig 문서를 참고하세요.

Update on GitHub

TRL

일반화 지식 증류 트레이너

개요

사용 팁

기대되는 데이터셋 유형

GKDTrainer

클래스 trl.experimental.gkd.GKDTrainer

train

save_model

push_to_hub

GKDConfig

클래스 trl.experimental.gkd.GKDConfig