# 汎用知識蒸留トレーナー

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

## 概要

汎化型知識蒸留（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](https://huggingface.co/papers/2306.13649) で提案されました。

論文の要旨は以下の通りです:

> Knowledge distillation (KD) は、より小さな student モデルを訓練することで teacher モデルを圧縮し、推論コストとメモリフットプリントを削減するために広く使用されています。しかし、現在の auto-regressive シーケンスモデル向けの KD 手法では、訓練中に参照される出力シーケンスと、推論時に student が生成する出力シーケンスとの間に分布の不一致が生じるという問題があります。この問題を解決するために、Generalized Knowledge Distillation (GKD) を提案します。GKD は、固定された出力シーケンスのセットだけに依存するのではなく、teacher からのフィードバックを活用して、student が自己生成した出力シーケンスで student を訓練します。教師あり KD のアプローチとは異なり、GKD は student と teacher の間で代替の損失関数を使用する柔軟性も提供します。これは、student が teacher の分布を模倣する表現力を持たない場合に有用です。さらに、GKD は蒸留と RL によるファインチューニング (RLHF) のシームレスな統合を可能にします。要約、翻訳、算術推論タスクにおける auto-regressive 言語モデルの蒸留、および instruction-tuning におけるタスク非依存の蒸留において、GKD の有効性を示します。

GKDの主な側面は以下の通りです：

1. 自己生成出力シーケンスを用いて学生モデルを訓練することで、自己回帰型シーケンスモデルにおける訓練と推論の分布不一致に対処します。
2. GKDは、汎化Jensen-Shannon分散（JSD）を通じて、教師モデルと生徒モデル間の異なる分散尺度を選択する柔軟性を提供します。これは、生徒モデルが教師モデルを完全に模倣する能力を欠いている場合に有用です。

このポストトレーニング手法は、[Kashif Rasul](https://huggingface.co/kashif)氏と[Lewis Tunstall](https://huggingface.co/lewtun)氏によって提供されました。

## 使用のコツ

[experimental.gkd.GKDTrainer](/docs/trl/main/ja/gkd_trainer#trl.experimental.gkd.GKDTrainer) は、教師モデルの引数を受け取る [SFTTrainer](/docs/trl/main/ja/sft_trainer#trl.SFTTrainer) クラスのラッパーです。[experimental.gkd.GKDConfig](/docs/trl/main/ja/gkd_trainer#trl.experimental.gkd.GKDConfig) 経由で設定する必要がある3つのパラメータは以下の通りです。

* `lmbda`: 学生データ比率、つまりオンポリシーの学生生成出力の割合を制御します。`lmbda=0.0`の場合、損失は教師のトークンレベルの確率で学生を訓練する教師付きJSDに帰着します。`lmbda=1.0`の場合、損失はオンポリシーJSDに帰着し、学生が出力シーケンスを生成し、教師からこれらのシーケンスに対するトークン固有のフィードバックを受けます。[0, 1]の間の値の場合、各バッチについて`lmbda`の値に基づいて、この2つのモードの間でランダムに選択されます。
* `seq_kd`: シーケンスレベルのKDを実行するかどうかを制御します（教師が生成した出力に対する教師ありFTと見なすことができます）。`seq_kd=True` かつ `lmbda=0.0` の場合、損失は教師ありJSDに帰着し、教師が出力シーケンスを生成し、学生は教師からこれらのシーケンスに対するトークンごとのフィードバックを受けます。
* `beta`: 汎化 Jensen-Shannon 分散における補間を制御します。`beta=0.0` の場合、損失は前方 KL 分散を近似し、`beta=1.0` の場合は後方 KL 分散を近似します。[0, 1] の間の値では、これら 2 つの間を補間します。

著者らは、オンポリシーデータ（`lmbda`が高い場合）の方がパフォーマンスが良く、最適な`beta`はタスクや評価方法によって異なると報告しています。

> [!WARNING]
> [Gemmaモデル](https://huggingface.co/models?other=gemma2)を学習する際は、`attn_implementation="kernels-community/flash-attn2"`が設定されていることを確認してください。そうしないと、このアーキテクチャが採用している[ソフトキャッピング技術](https://huggingface.co/blog/gemma2#soft-capping-and-attention-implementations)により、logitsにNaNが発生します。

基本的なAPIは以下の通りです:

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

### 期待されるデータセットの型

データセットは、各メッセージが以下のキーを持つ辞書のリストである「messages」のリストとしてフォーマットされる必要があります。

* `role`: `system`、`assistant`、または `user`
* `content`: メッセージの内容

## 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)
```

[ソース](https://github.com/huggingface/trl/blob/main/trl/experimental/gkd/gkd_trainer.py#L54)

**パラメータ:**

model ([PreTrainedModel](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel) または `torch.nn.Module` または `str`、*optional*) : 学習対象のモデル、または事前学習済みモデルからインスタンス化されるモデルの文字列識別子。

teacher_model ([PreTrainedModel](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel) または `torch.nn.Module` または `str`、*optional*) : 知識蒸留用の教師モデル、または事前学習済みモデルからインスタンス化されるモデルの文字列識別子。

args ([experimental.gkd.GKDConfig](/docs/trl/main/ja/gkd_trainer#trl.experimental.gkd.GKDConfig)、*任意*) : 学習引数。

data_collator (`DataCollator`, *任意*) : データセットからサンプルをバッチ化するデータコラレーターです。デフォルトでは、`processing_class` を使用した `experimental.utils.DataCollatorForChatML` が使用されます。

train_dataset ([Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset), *optional*) : 学習用のデータセット。

eval_dataset ([Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset) または [Dataset](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset) の `dict`、*optional*) : 評価用のデータセット。

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) または [ProcessorMixin](https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin)、*任意*) : データを処理するクラス。

compute_metrics (`Callable`, *任意*) : 評価時にメトリクスを計算する関数。[EvalPrediction](https://huggingface.co/docs/transformers/main/en/internal/trainer_utils#transformers.EvalPrediction) を受け取り、文字列から float への辞書を返す必要があります。

callbacks (`list` of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback), *optional*) : 学習中に使用するコールバック。

optimizers (`tuple` of `torch.optim.Optimizer` and `torch.optim.lr_scheduler.LambdaLR`, *optional*, defaults to `(None, None)`) : 学習に使用するオプティマイザーと学習率スケジューラーを含むタプル。

preprocess_logits_for_metrics (`Callable`, *任意*) : メトリクスの計算前に logits を前処理する関数。`logits` と `labels` を受け取り、メトリクスの計算に使用する logits を返す必要があります。

peft_config ([PeftConfig](https://huggingface.co/docs/peft/main/en/package_reference/config#peft.PeftConfig), *任意*) : PEFT をトレーニングに使用するための PEFT 設定。`None` の場合、PEFT は使用されません。指定された場合、`model` は指定された PEFT アダプターでラップされます。

formatting_func (`Callable`, *任意*) : データセットを整形する関数。引数として1つのサンプルを受け取り、サンプルを返す必要があります。

言語モデルの汎化知識蒸留（GKD）用のTrainer。

For details on GKD, see the paper: [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)
```

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L1447)

**パラメータ:**

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[[trl.experimental.gkd.GKDTrainer.save_model]]

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

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4036)

モデルを保存するため、`from_pretrained()`を使用して再読み込みできます。

メインプロセスからのみ保存します。

#### 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)
```

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4307)

**パラメータ:**

commit_message (`str`、*任意*、デフォルトは`"End of training"`) : push 時にコミットするメッセージ。

blocking (`bool`、*任意*、デフォルトは`True`) : `git push`が完了した後にのみ関数が返すかどうか。

token (`str`、*任意*、デフォルトは`None`) : Trainer の元の引数を上書きするための、書き込み権限を持つトークン。

revision (`str`, *任意*) : コミット元の git リビジョンです。デフォルトでは "main" ブランチの先頭です。

kwargs (`dict[str, Any]`、*省略可能*) : `~Trainer.create_model_card` に渡される追加のキーワード引数。

**戻り値:**

`blocking=False`の場合、モデルがプッシュされたリポジトリのURL、または`blocking=True`の場合、コミットの進捗を追跡する`Future`オブジェクトを返します。

`self.args.hub_model_id` のリポジトリに、`self.model` と `self.processing_class` を 🤗 model hub にアップロードします。

## 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)
```

[ソース](https://github.com/huggingface/trl/blob/main/trl/experimental/gkd/gkd_config.py#L22)

**パラメータ:**

temperature (`float`、*任意*、デフォルトは`0.9`) : サンプリング用の温度です。温度が高いほど、生成される補完はランダムになります。

lmbda (`float`、*任意*、デフォルトは`0.5`) : 生徒データの割合（つまり、オンポリシーで生徒が生成した出力の割合）を制御するラムダパラメータ。

beta (`float`、*optional*、デフォルトは`0.5`) : Generalized Jensen-Shannon Divergence損失における`0.0`と`1.0`の間の補間係数。betaが`0.0`の場合、損失はKL分散になります。betaが`1.0`の場合、損失はInverse KL Divergenceになります。

max_new_tokens (`int`、*optional*、デフォルトは`128`) : 1回の補完で生成するトークンの最大数。

teacher_model_name_or_path (`str`, *任意*) : 教師モデルの名前またはパス。`None`の場合、教師モデルは学習対象のモデルと同じになります。

teacher_model_init_kwargs (`dict[str, Any]`, *省略可能*) : 文字列から教師モデルをインスタンス化する場合に、`AutoModelForCausalLM.from_pretrained` に渡すキーワード引数。

disable_dropout (`bool`、*任意*、デフォルトは`True`) : モデルの dropout を無効にするかどうか。

seq_kd (`bool`、*任意*、デフォルトは`False`) : Seq_kd パラメータ。Sequence-Level KD を実行するかどうかを制御します（教師モデルが生成した出力に対する教師あり FT と見なすことができます）。

[experimental.gkd.GKDTrainer](/docs/trl/main/ja/gkd_trainer#trl.experimental.gkd.GKDTrainer)の構成クラスです。

このクラスには、GKD 学習に固有のパラメータのみが含まれています。学習引数の完全なリストについては、[TrainingArguments](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments) および [SFTConfig](/docs/trl/main/ja/sft_trainer#trl.SFTConfig) のドキュメントを参照してください。