Hugging Face's logo

TRL documentation

Pelatih Distilasi Pengetahuan yang Digeneralisasi

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

Pelatih Distilasi Pengetahuan yang Digeneralisasi

model badge

Ikhtisar

Generalized Knowledge Distillation (GKD) diusulkan dalam On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes oleh Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, dan Olivier Bachem.

Abstrak dari makalah tersebut adalah sebagai berikut:

Knowledge distillation (KD) banyak digunakan untuk mengompresi model guru guna mengurangi biaya inferensi dan jejak memorinya, dengan melatih model siswa yang lebih kecil. Namun, metode KD saat ini untuk model sekuens auto-regresif mengalami ketidakcocokan distribusi antara sekuens output yang dilihat selama pelatihan dan yang dihasilkan oleh siswa selama inferensi. Untuk mengatasi masalah ini, kami memperkenalkan Generalized Knowledge Distillation (GKD). Alih-alih hanya mengandalkan set sekuens output yang tetap, GKD melatih siswa pada sekuens output yang dihasilkan sendiri oleh siswa dengan memanfaatkan umpan balik dari guru pada sekuens tersebut. Berbeda dengan pendekatan KD yang diawasi, GKD juga menawarkan fleksibilitas untuk menggunakan fungsi loss alternatif antara siswa dan guru, yang dapat berguna ketika siswa tidak memiliki ekspresivitas untuk meniru distribusi guru. Selain itu, GKD memfasilitasi integrasi mulus antara distilasi dengan penyesuaian halus RL (RLHF). Kami mendemonstrasikan efektivitas GKD untuk mendistilasi model bahasa auto-regresif pada tugas ringkasan, terjemahan, dan penalaran aritmatika, serta distilasi yang tidak bergantung pada tugas untuk penyetelan instruksi.

Aspek-aspek utama dari GKD adalah:

  1. Hal ini mengatasi ketidakcocokan distribusi antara pelatihan dan inferensi pada model sekuens auto-regresif dengan melatih model siswa pada sekuens output yang dihasilkan oleh dirinya sendiri.
  2. GKD memberikan fleksibilitas dalam memilih berbagai ukuran divergensi antara model siswa dan model guru melalui Generalized Jensen-Shannon Divergence (JSD), yang dapat berguna ketika model siswa tidak memiliki kapasitas untuk sepenuhnya meniru model guru.

Metode pasca-pelatihan ini disumbangkan oleh Kashif Rasul dan Lewis Tunstall.

Tips penggunaan

experimental.gkd.GKDTrainer adalah pembungkus di sekitar kelas SFTTrainer yang menerima argumen model guru. Ia membutuhkan tiga parameter untuk diatur melalui experimental.gkd.GKDConfig, yaitu:

  • lmbda: mengontrol fraksi data student, yaitu proporsi output yang dihasilkan student secara on-policy. Ketika lmbda=0.0, loss berkurang menjadi JSD supervised di mana student dilatih dengan probabilitas tingkat token dari teacher. Ketika lmbda=1.0, loss berkurang menjadi JSD on-policy, di mana student menghasilkan urutan output dan umpan balik spesifik token pada urutan tersebut dari teacher. Untuk nilai di antara [0, 1], prosesnya acak di antara keduanya berdasarkan nilai lmbda untuk setiap batch.
  • seq_kd: mengontrol apakah akan melakukan Sequence-Level KD (dapat dipandang sebagai supervised FT pada output yang dihasilkan teacher). Ketika seq_kd=True dan lmbda=0.0, loss berkurang menjadi supervised JSD, di mana teacher menghasilkan sekuens output dan student menerima umpan balik spesifik token pada sekuens tersebut dari teacher.
  • beta: mengontrol interpolasi dalam Divergensi Jensen-Shannon yang digeneralisasi. Ketika beta=0.0, loss mendekati divergensi KL forward, sedangkan untuk beta=1.0, loss mendekati divergensi KL reverse. Untuk nilai di antara [0, 1], loss melakukan interpolasi di antara keduanya.

Para penulis menemukan bahwa data on-policy (dengan lmbda tinggi) memberikan hasil yang lebih baik, dan nilai optimal beta bervariasi tergantung pada tugas dan metode evaluasi.

Pastikan attn_implementation="kernels-community/flash-attn2" saat melatih model Gemma. Jika tidak, Anda akan menemui nilai NaN pada logits karena teknik soft capping yang diadopsi oleh arsitektur ini.

API dasarnya adalah sebagai berikut:

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

Jenis dataset yang diharapkan

Dataset harus diformat sebagai daftar “messages” di mana setiap message adalah daftar kamus dengan kunci-kunci berikut:

  • role: salah satu dari system, assistant, atau user
  • content: isi pesan

GKDTrainer

kelas trl.experimental.gkd.GKDTrainer

< >

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

Parameter

  • model (PreTrainedModel atau torch.nn.Module atau str, opsional) — Model yang akan dilatih, atau pengidentifikasi string dari model yang akan diinstansiasi dari model pretrained.
  • teacher_model (PreTrainedModel atau torch.nn.Module atau str, opsional) — Model guru untuk knowledge distillation, atau pengidentifikasi string dari model yang akan diinstansiasi dari model pretrained.
  • args (experimental.gkd.GKDConfig, opsional) — Argumen pelatihan.
  • data_collator (DataCollator, opsional) — Data collator untuk membuat batch sampel dari dataset. Nilainya secara default adalah experimental.utils.DataCollatorForChatML yang menggunakan processing_class.
  • train_dataset (Dataset, opsional) — Dataset untuk pelatihan.
  • eval_dataset (Dataset atau dict dari Dataset, opsional) — Dataset untuk evaluasi.
  • processing_class (PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin atau ProcessorMixin, opsional) — Kelas untuk memproses data.
  • compute_metrics (Callable, opsional) — Fungsi untuk menghitung metrik saat evaluasi. Harus menerima EvalPrediction dan mengembalikan dictionary string ke float.
  • callbacks (list dari TrainerCallback, opsional) — Callback yang digunakan selama pelatihan.
  • optimizers (tuple dari torch.optim.Optimizer dan torch.optim.lr_scheduler.LambdaLR, opsional, default (None, None)) — Tuple yang berisi optimizer dan learning rate scheduler yang akan digunakan untuk pelatihan.
  • preprocess_logits_for_metrics (Callable, opsional) — Fungsi untuk memproses ulang logits sebelum menghitung metrik. Harus menerima logits dan labels serta mengembalikan logits yang akan digunakan untuk perhitungan metrik.
  • peft_config (PeftConfig, opsional) — Konfigurasi PEFT untuk menggunakan PEFT dalam pelatihan. Jika None, PEFT tidak digunakan. Jika disediakan, model akan dibungkus dengan adapter PEFT yang ditentukan.
  • formatting_func (Callable, opsional) — Fungsi untuk memformat dataset. Harus menerima satu contoh dan mengembalikan satu contoh.

Trainer untuk Generalized Knowledge Distillation (GKD) dari model bahasa.

Untuk detail tentang GKD, lihat makalah: On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.

latih

< >

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

Parameter

  • resume_from_checkpoint (str atau bool, opsional) — Jika berupa str, jalur lokal ke checkpoint yang disimpan oleh instance Trainer sebelumnya. Jika berupa bool dan bernilai True, muat checkpoint terakhir di args.output_dir yang disimpan oleh instance Trainer sebelumnya. Jika ada, pelatihan akan dilanjutkan dari status model/optimizer/scheduler yang dimuat di sini.
  • trial (optuna.Trial atau dict[str, Any], opsional) — Jalankan trial atau kamus hiperparameter untuk pencarian hiperparameter.
  • ignore_keys_for_eval (list[str], opsional) — Daftar kunci dalam output model Anda (jika berupa dictionary) yang harus diabaikan saat mengumpulkan prediksi untuk evaluasi selama pelatihan.

Mengembalikan

~trainer_utils.TrainOutput

Objek yang berisi jumlah langkah global, kerugian pelatihan, dan metrik.

Titik masuk utama pelatihan.

save_model

< >

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

Akan menyimpan model, sehingga Anda dapat memuatnya kembali menggunakan from_pretrained().

Hanya akan menyimpan dari proses utama.

push_to_hub

< >

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

Parameter

  • commit_message (str, opsional, default "End of training") — Pesan untuk commit saat melakukan push.
  • blocking (bool, opsional, default True) — Apakah fungsi harus hanya mengembalikan nilai setelah git push selesai.
  • token (str, opsional, defaultnya None) — Token dengan izin tulis untuk menimpa argumen asli Trainer.
  • revision (str, opsional) — Revisi git untuk di-commit. Secara default menggunakan head dari branch “main”.
  • kwargs (dict[str, Any], opsional) — Argumen kata kunci tambahan yang diteruskan ke ~Trainer.create_model_card.

Unggah self.model dan self.processing_class ke 🤗 model hub di repo self.args.hub_model_id.

GKDConfig

kelas trl.experimental.gkd.GKDConfig

< >

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

Parameter

  • temperature (float, opsional, default 0.9) — Suhu untuk sampling. Semakin tinggi suhu, semakin acak hasil pelengkapannya.
  • lmbda (float, opsional, default 0.5) — Parameter Lambda yang mengontrol fraksi data student (yaitu, proporsi output yang dihasilkan student secara on-policy).
  • beta (float, opsional, default 0.5) — Koefisien interpolasi antara 0.0 dan 1.0 untuk loss Generalized Jensen-Shannon Divergence. Ketika beta adalah 0.0, loss-nya adalah KL divergence. Ketika beta adalah 1.0, loss-nya adalah Inverse KL Divergence.
  • max_new_tokens (int, opsional, defaultnya 128) — Jumlah maksimum token yang akan dihasilkan per penyelesaian.
  • teacher_model_name_or_path (str, opsional) — Nama atau path dari model teacher. Jika None, model teacher akan sama dengan model yang sedang dilatih.
  • teacher_model_init_kwargs (dict[str, Any], opsional) — Argumen kata kunci untuk diteruskan ke AutoModelForCausalLM.from_pretrained saat menginstansiasi model teacher dari string.
  • disable_dropout (bool, opsional, default True) — Apakah akan menonaktifkan dropout pada model.
  • seq_kd (bool, opsional, default False) — Parameter seq_kd yang mengontrol apakah akan melakukan Sequence-Level KD (dapat dipandang sebagai supervised FT pada output yang dihasilkan teacher).

Kelas konfigurasi untuk experimental.gkd.GKDTrainer.

Kelas ini hanya mencakup parameter yang spesifik untuk pelatihan GKD. Untuk daftar lengkap argumen pelatihan, silakan merujuk ke dokumentasi TrainingArguments dan SFTConfig.

Update on GitHub

←CPO GMPO→