# Trainer Chưng cất Kiến thức Tổng quát

[![model badge](https://img.shields.io/badge/All_models-GKD-blue)](https://huggingface.co/models?other=gkd,trl)

## Tổng quan

Generalized Knowledge Distillation (GKD) được đề xuất trong [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://huggingface.co/papers/2306.13649) bởi Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist và Olivier Bachem.

Tóm tắt từ bài báo như sau:

> Knowledge distillation (KD) được sử dụng rộng rãi để nén một mô hình teacher nhằm giảm chi phí suy luận và dung lượng bộ nhớ, bằng cách huấn luyện một mô hình student nhỏ hơn. Tuy nhiên, các phương pháp KD hiện tại cho các mô hình chuỗi tự hồi quy (auto-regressive) gặp phải vấn đề lệch phân phối giữa các chuỗi đầu ra được xem trong quá trình huấn luyện và các chuỗi do student tạo ra trong quá trình suy luận. Để giải quyết vấn đề này, chúng tôi giới thiệu Generalized Knowledge Distillation (GKD). Thay vì chỉ dựa vào một tập hợp cố định các chuỗi đầu ra, GKD huấn luyện student trên các chuỗi đầu ra do chính nó tạo ra bằng cách tận dụng phản hồi từ teacher trên các chuỗi đó. Khác với các tiếp cận KD có giám sát, GKD còn cung cấp sự linh hoạt để sử dụng các hàm mất mát (loss function) thay thế giữa student và teacher, điều này có thể hữu ích khi student thiếu khả năng diễn đạt (expressivity) để bắt chước phân phối của teacher. Hơn nữa, GKD tạo điều kiện tích hợp liền mạch quá trình distillation với tinh chỉnh bằng RL (RLHF). Chúng tôi chứng minh hiệu quả của GKD trong việc distill các mô hình ngôn ngữ tự hồi quy trên các tác vụ tóm tắt, dịch thuật và suy luận toán học, cũng như distillation không phụ thuộc tác vụ cho việc instruction-tuning.

Các khía cạnh chính của GKD là:

1. Nó giải quyết sự lệch phân phối giữa huấn luyện và suy luận trong các mô hình chuỗi tự hồi quy bằng cách huấn luyện mô hình học sinh trên các chuỗi đầu ra do chính nó tạo ra.
2. GKD cho phép linh hoạt trong việc lựa chọn các phép đo phân kỳ khác nhau giữa mô hình học sinh và mô hình giáo viên thông qua phân kỳ Jensen-Shannon tổng quát (JSD), điều này có thể hữu ích khi mô hình học sinh không đủ khả năng để mô phỏng hoàn toàn mô hình giáo viên.

Phương pháp hậu huấn luyện này được đóng góp bởi [Kashif Rasul](https://huggingface.co/kashif) và [Lewis Tunstall](https://huggingface.co/lewtun).

## Mẹo sử dụng

[experimental.gkd.GKDTrainer](/docs/trl/main/vi/gkd_trainer#trl.experimental.gkd.GKDTrainer) là một lớp bao quanh lớp [SFTTrainer](/docs/trl/main/vi/sft_trainer#trl.SFTTrainer), nhận một đối số mô hình giáo viên. Nó cần ba tham số được đặt thông qua [experimental.gkd.GKDConfig](/docs/trl/main/vi/gkd_trainer#trl.experimental.gkd.GKDConfig), cụ thể là:

* `lmbda`:  kiểm soát tỷ lệ dữ liệu của student, tức là phần trăm các đầu ra do student tạo ra theo chính sách (on-policy). Khi `lmbda=0.0`, loss thu gọn về JSD có giám sát, trong đó student được huấn luyện với xác suất ở cấp độ token của teacher. Khi `lmbda=1.0`, loss thu gọn về JSD on-policy, trong đó student tạo ra các chuỗi đầu ra và nhận phản hồi cụ thể cho từng token trên các chuỗi này từ teacher. Với các giá trị nằm trong khoảng [0, 1], quá trình sẽ ngẫu nhiên giữa hai chế độ dựa trên giá trị `lmbda` cho mỗi batch.
* `seq_kd`:  kiểm soát việc có thực hiện Sequence-Level KD hay không (có thể xem là supervised FT trên out do teacher sinh ra). Khi `seq_kd=True` và `lmbda=0.0`, loss sẽ rút gọn về supervised JSD, trong đó teacher sinh ra các chuỗi output và student nhận phản hồi cụ thể cho từng token trên các chuỗi này từ teacher.
* `beta`: kiểm soát sự nội suy trong phân kỳ Jensen-Shannon tổng quát. Khi `beta=0.0`, loss xấp xỉ phân kỳ KL thuận, trong khi với `beta=1.0`, loss xấp xỉ phân kỳ KL nghịch. Với các giá trị nằm trong khoảng [0, 1], nó nội suy giữa hai phân kỳ này.

Các tác giả nhận thấy rằng dữ liệu on-policy (giá trị `lmbda` cao) cho kết quả tốt hơn và giá trị tối ưu của `beta` thay đổi tùy thuộc vào tác vụ và phương pháp đánh giá.

> [!WARNING]
> Hãy đảm bảo rằng `attn_implementation="kernels-community/flash-attn2"` khi huấn luyện [mô hình Gemma](https://huggingface.co/models?other=gemma2). Nếu không, bạn sẽ gặp các giá trị NaN trong logits do [kỹ thuật soft capping](https://huggingface.co/blog/gemma2#soft-capping-and-attention-implementations) được áp dụng bởi kiến trúc này.

API cơ bản như sau:

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

### Kiểu dữ liệu tập dữ liệu dự kiến

Bộ dữ liệu nên được định dạng dưới dạng một danh sách các "messages", trong đó mỗi message là một danh sách các từ điển có các khóa sau:

* `role`: là `system`, `assistant` hoặc `user`
* `content`: nội dung tin nhắn

## GKDTrainer[[trl.experimental.gkd.GKDTrainer]]

#### trl.experimental.gkd.GKDTrainer[[trl.experimental.gkd.GKDTrainer]]

```python
trl.experimental.gkd.GKDTrainer(model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str, NoneType] = None, teacher_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str] = None, args: trl.experimental.gkd.gkd_config.GKDConfig | None = None, data_collator: collections.abc.Callable[[list[typing.Any]], dict[str, typing.Any]] | None = None, train_dataset: datasets.arrow_dataset.Dataset | None = None, eval_dataset: datasets.arrow_dataset.Dataset | dict[str, datasets.arrow_dataset.Dataset] | None = None, processing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.image_processing_utils.BaseImageProcessor | transformers.feature_extraction_utils.FeatureExtractionMixin | transformers.processing_utils.ProcessorMixin | None = None, compute_metrics: collections.abc.Callable[[transformers.trainer_utils.EvalPrediction], dict] | None = None, callbacks: list[transformers.trainer_callback.TrainerCallback] | None = None, optimizers: 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 = None, peft_config: PeftConfig | None = None, formatting_func: collections.abc.Callable | None = None)
```

[Nguồn](https://github.com/huggingface/trl/blob/main/trl/experimental/gkd/gkd_trainer.py#L54)

**Tham số:**

model ([PreTrainedModel](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel) hoặc `torch.nn.Module` hoặc `str`, *optional*) : Mô hình cần được huấn luyện, hoặc chuỗi định danh của mô hình cần được khởi tạo từ một mô hình được huấn luyện trước.

teacher_model ([PreTrainedModel](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel) hoặc `torch.nn.Module` hoặc `str`, *optional*) : Mô hình teacher cho knowledge distillation, hoặc chuỗi định danh của mô hình cần được khởi tạo từ một mô hình đã được huấn luyện trước.

args ([experimental.gkd.GKDConfig](/docs/trl/main/vi/gkd_trainer#trl.experimental.gkd.GKDConfig), *tùy chọn*) : Các đối số huấn luyện.

data_collator (`DataCollator`, *tùy chọn*) : Data collator để nhóm các mẫu từ tập dữ liệu. Mặc định là `experimental.utils.DataCollatorForChatML` sử dụng `processing_class`.

train_dataset ([Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset), *tùy chọn*) : Dataset để huấn luyện.

eval_dataset ([Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset) hoặc `dict` của [Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset), *tuỳ chọn*) : Dataset để đánh giá.

processing_class ([PreTrainedTokenizerBase](https://huggingface.co/docs/transformers/main/en/internal/tokenization_utils#transformers.PreTrainedTokenizerBase), [BaseImageProcessor](https://huggingface.co/docs/transformers/main/en/main_classes/image_processor#transformers.BaseImageProcessor), [FeatureExtractionMixin](https://huggingface.co/docs/transformers/main/en/main_classes/feature_extractor#transformers.FeatureExtractionMixin) hoặc [ProcessorMixin](https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin), *tùy chọn*) : Lớp để xử lý dữ liệu.

compute_metrics (`Callable`, *tùy chọn*) : Hàm để tính các chỉ số đánh giá. Phải nhận vào một [EvalPrediction](https://huggingface.co/docs/transformers/main/en/internal/trainer_utils#transformers.EvalPrediction) và trả về một từ điển chuỗi sang số float.

callbacks (`list` của [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback), *optional*) : Các callback được sử dụng trong quá trình huấn luyện.

optimizers (`tuple` của `torch.optim.Optimizer` và `torch.optim.lr_scheduler.LambdaLR`, *không bắt buộc*, mặc định là `(None, None)`) : Tuple chứa optimizer và learning rate scheduler sẽ được sử dụng cho quá trình huấn luyện.

preprocess_logits_for_metrics (`Callable`, *tùy chọn*) : Hàm để tiền xử lý logits trước khi tính toán các chỉ số. Phải nhận vào `logits` và `labels` và trả về logits sẽ được sử dụng để tính toán chỉ số.

peft_config ([PeftConfig](https://huggingface.co/docs/peft/main/en/package_reference/config#peft.PeftConfig), *tùy chọn*) : Cấu hình PEFT để sử dụng PEFT cho việc huấn luyện. Nếu là `None`, PEFT sẽ không được sử dụng. Nếu được cung cấp, `model` sẽ được bọc với bộ thích ứng PEFT được chỉ định.

formatting_func (`Callable`, *tùy chọn*) : Hàm để định dạng bộ dữ liệu. Phải nhận vào một ví dụ và trả về một ví dụ.

Trainer cho Chưng cất Kiến thức Tổng quát (GKD) của các mô hình ngôn ngữ.

Để biết thêm chi tiết về GKD, hãy xem bài báo: [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://huggingface.co/papers/2306.13649).

#### train[[trl.experimental.gkd.GKDTrainer.train]]

```python
train(resume_from_checkpoint: str | bool | None = None, trial: optuna.Trial | dict[str, Any] | None = None, ignore_keys_for_eval: list[str] | None = None)
```

[Nguồn](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L1447)

**Tham số:**

resume_from_checkpoint (`str` hoặc `bool`, *tùy chọn*) : Nếu là `str`, đường dẫn cục bộ đến checkpoint đã lưu bởi một phiên `Trainer` trước đó. Nếu là `bool` và bằng `True`, nạp checkpoint cuối cùng trong *args.output_dir* do một phiên `Trainer` trước đó lưu lại. Nếu có, quá trình huấn luyện sẽ tiếp tục từ trạng thái model/optimizer/scheduler được nạp tại đây.

trial (`optuna.Trial` hoặc `dict[str, Any]`, *tùy chọn*) : Chạy thử nghiệm hoặc từ điển siêu tham số cho việc tìm kiếm siêu tham số.

ignore_keys_for_eval (`list[str]`, *tùy chọn*) : Một danh sách các khóa trong đầu ra của mô hình của bạn (nếu nó là một từ điển) cần được bỏ qua khi thu thập các dự đoán để đánh giá trong quá trình huấn luyện.

**Trả về:** `~trainer_utils.TrainOutput`

Đối tượng chứa số bước toàn cục, tổn thất huấn luyện và các chỉ số.

Điểm vào chính cho quá trình huấn luyện.

#### save_model[[trl.experimental.gkd.GKDTrainer.save_model]]

```python
save_model(output_dir: str | None = None, _internal_call: bool = False)
```

[Nguồn](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4036)

Sẽ lưu mô hình, để bạn có thể tải lại bằng `from_pretrained()`.

Chỉ lưu từ tiến trình chính.

#### push_to_hub[[trl.experimental.gkd.GKDTrainer.push_to_hub]]

```python
push_to_hub(commit_message: str | None = 'End of training', blocking: bool = True, token: str | None = None, revision: str | None = None, **kwargs)
```

[Nguồn](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4307)

**Tham số:**

commit_message (`str`, *tùy chọn*, mặc định là `"End of training"`) : Thông điệp để commit khi đẩy lên.

blocking (`bool`, *tùy chọn*, mặc định là `True`) : Xác định xem hàm có nên chỉ trả về khi `git push` đã hoàn tất hay không.

token (`str`, *tùy chọn*, mặc định là `None`) : Token có quyền ghi để ghi đè lên các đối số gốc của Trainer.

revision (`str`, *tùy chọn*) : Bản sửa git để commit từ. Mặc định là đầu của nhánh "main".

kwargs (`dict[str, Any]`, *tùy chọn*) : Các đối số từ khóa bổ sung được truyền đến `~Trainer.create_model_card`.

**Trả về:**

URL của kho lưu trữ nơi mô hình đã được đẩy lên nếu `blocking=False`, hoặc một đối tượng `Future` theo dõi tiến trình của lệnh commit nếu `blocking=True`.

Tải `self.model` và `self.processing_class` lên 🤗 model hub tại repo `self.args.hub_model_id`.

## GKDConfig[[trl.experimental.gkd.GKDConfig]]

#### trl.experimental.gkd.GKDConfig[[trl.experimental.gkd.GKDConfig]]

```python
trl.experimental.gkd.GKDConfig(output_dir: str | None = None, per_device_train_batch_size: int = 8, num_train_epochs: float = 3.0, max_steps: int = -1, learning_rate: float = 2e-05, lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear', lr_scheduler_kwargs: dict | str | None = None, warmup_steps: float = 0, optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused', optim_args: str | None = None, weight_decay: float = 0.0, adam_beta1: float = 0.9, adam_beta2: float = 0.999, adam_epsilon: float = 1e-08, optim_target_modules: None | str | list[str] = None, gradient_accumulation_steps: int = 1, average_tokens_across_devices: bool = True, max_grad_norm: float = 1.0, label_smoothing_factor: float = 0.0, bf16: bool | None = None, fp16: bool = False, bf16_full_eval: bool = False, fp16_full_eval: bool = False, tf32: bool | None = None, gradient_checkpointing: bool = True, gradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = None, torch_compile: bool = False, torch_compile_backend: str | None = None, torch_compile_mode: str | None = None, use_liger_kernel: bool = False, liger_kernel_config: dict[str, bool] | None = None, use_cache: bool = False, neftune_noise_alpha: float | None = None, torch_empty_cache_steps: int | None = None, auto_find_batch_size: bool = False, logging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps', logging_steps: float = 10, logging_first_step: bool = False, log_on_each_node: bool = True, logging_nan_inf_filter: bool = True, include_num_input_tokens_seen: str | bool = 'no', log_level: str = 'passive', log_level_replica: str = 'warning', disable_tqdm: bool | None = None, report_to: None | str | list[str] = 'none', run_name: str | None = None, project: str = 'huggingface', trackio_space_id: str | None = None, trackio_bucket_id: str | None = None, trackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = None, eval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no', eval_steps: float | None = None, eval_delay: float = 0, per_device_eval_batch_size: int = 8, prediction_loss_only: bool = False, eval_on_start: bool = False, eval_do_concat_batches: bool = True, eval_use_gather_object: bool = False, eval_accumulation_steps: int | None = None, include_for_metrics: list[str] = <factory>, batch_eval_metrics: bool = False, save_only_model: bool = False, save_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps', save_steps: float = 500, save_on_each_node: bool = False, save_total_limit: int | None = None, enable_jit_checkpoint: bool = False, push_to_hub: bool = False, hub_token: str | None = None, hub_private_repo: bool | None = None, hub_model_id: str | None = None, hub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save', hub_always_push: bool = False, hub_revision: str | None = None, load_best_model_at_end: bool = False, metric_for_best_model: str | None = None, greater_is_better: bool | None = None, ignore_data_skip: bool = False, restore_callback_states_from_checkpoint: bool = False, full_determinism: bool = False, seed: int = 42, data_seed: int | None = None, use_cpu: bool = False, accelerator_config: dict | str | None = None, parallelism_config: accelerate.parallelism_config.ParallelismConfig | None = None, dataloader_drop_last: bool = False, dataloader_num_workers: int = 0, dataloader_pin_memory: bool = True, dataloader_persistent_workers: bool = False, dataloader_prefetch_factor: int | None = None, dataloader_multiprocessing_context: str | None = None, dataloader_in_order: bool = True, remove_unused_columns: bool = True, label_names: list[str] | None = None, train_sampling_strategy: str = 'random', length_column_name: str = 'length', ddp_find_unused_parameters: bool | None = None, ddp_bucket_cap_mb: int | None = None, ddp_broadcast_buffers: bool | None = None, ddp_static_graph: bool | None = None, ddp_backend: str | None = None, ddp_timeout: int = 1800, fsdp: str | None = None, fsdp_config: dict[str, typing.Any] | str | None = None, deepspeed: dict | str | None = None, debug: str | list[transformers.debug_utils.DebugOption] = '', skip_memory_metrics: bool = True, do_train: bool = False, do_eval: bool = False, do_predict: bool = False, resume_from_checkpoint: str | None = None, local_rank: int = -1, model_init_kwargs: dict[str, typing.Any] | str | None = None, trust_remote_code: bool = False, router_aux_loss_coef: float | None = None, chat_template_path: str | None = None, dataset_text_field: str = 'text', dataset_kwargs: dict[str, typing.Any] | str | None = None, dataset_num_proc: int | None = None, eos_token: str | None = None, max_length: int | None = 1024, truncation_mode: str = 'keep_start', shuffle_dataset: bool = False, packing: bool = False, packing_strategy: str = 'bfd', padding_free: bool = False, pad_to_multiple_of: int | str | None = 'auto', eval_packing: bool | None = None, completion_only_loss: bool | None = None, assistant_only_loss: bool = False, loss_type: str = 'nll', activation_offloading: bool = False, pad_token: str | None = None, temperature: float = 0.9, lmbda: float = 0.5, beta: float = 0.5, max_new_tokens: int = 128, teacher_model_name_or_path: str | None = None, teacher_model_init_kwargs: dict[str, typing.Any] | str | None = None, disable_dropout: bool = True, seq_kd: bool = False)
```

[Nguồn](https://github.com/huggingface/trl/blob/main/trl/experimental/gkd/gkd_config.py#L22)

**Tham số:**

temperature (`float`, *tùy chọn*, mặc định là `0.9`) : Nhiệt độ cho việc lấy mẫu. Nhiệt độ càng cao, các phần hoàn thành càng ngẫu nhiên hơn.

lmbda (`float`, *tuỳ chọn*, mặc định là `0.5`) : Tham số Lambda kiểm soát tỷ lệ dữ liệu của học sinh (tức là tỷ lệ các đầu ra do học sinh tạo ra theo chính sách).

beta (`float`, *optional*, mặc định là `0.5`) : Hệ số nội suy giữa `0.0` và `1.0` của hàm mất mát Generalized Jensen-Shannon Divergence. Khi beta là `0.0`, hàm mất mát là KL divergence. Khi beta là `1.0`, hàm mất mát là Inverse KL Divergence.

max_new_tokens (`int`, *tùy chọn*, mặc định là `128`) : Số lượng token tối đa để tạo ra cho mỗi lần hoàn thành.

teacher_model_name_or_path (`str`, *tùy chọn*) : Tên hoặc đường dẫn của mô hình teacher. Nếu là `None`, mô hình teacher sẽ giống với mô hình đang được huấn luyện.

teacher_model_init_kwargs (`dict[str, Any]`, *tùy chọn*) : Các đối số từ khóa để truyền cho `AutoModelForCausalLM.from_pretrained` khi khởi tạo mô hình teacher từ một chuỗi.

disable_dropout (`bool`, *tùy chọn*, mặc định là `True`) : Xác định có vô hiệu hóa dropout trong mô hình hay không.

seq_kd (`bool`, *tùy chọn*, mặc định là `False`) : Tham số Seq_kd kiểm soát việc có thực hiện Sequence-Level KD (có thể xem là FT có giám sát trên đầu ra do teacher tạo ra) hay không.

Lớp cấu hình cho [experimental.gkd.GKDTrainer](/docs/trl/main/vi/gkd_trainer#trl.experimental.gkd.GKDTrainer).

Lớp này chỉ bao gồm các tham số đặc thù cho việc huấn luyện GKD. Để xem danh sách đầy đủ các đối số huấn luyện, vui lòng tham khảo tài liệu [TrainingArguments](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments) và [SFTConfig](/docs/trl/main/vi/sft_trainer#trl.SFTConfig).

