汎用知識蒸留トレーナー

概要

汎化型知識蒸留（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 モデルを訓練することで 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の主な側面は以下の通りです：

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

このポストトレーニング手法は、Kashif Rasul氏とLewis Tunstall氏によって提供されました。

使用のコツ

experimental.gkd.GKDTrainer は、教師モデルの引数を受け取る SFTTrainer クラスのラッパーです。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はタスクや評価方法によって異なると報告しています。

Gemmaモデルを学習する際は、attn_implementation="kernels-community/flash-attn2"が設定されていることを確認してください。そうしないと、このアーキテクチャが採用しているソフトキャッピング技術により、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、optional) — 学習対象のモデル、または事前学習済みモデルからインスタンス化されるモデルの文字列識別子。
teacher_model (PreTrainedModel または torch.nn.Module または str、任意) — 知識蒸留用の教師モデル、または事前学習モデルからインスタンス化されるモデルの文字列識別子。
args (experimental.gkd.GKDConfig、任意) — 学習引数。
data_collator (DataCollator, 省略可) — データセットからサンプルをバッチ化する際に使用するデータコラレーター。デフォルトでは、processing_class を使用した experimental.utils.DataCollatorForChatML が使用されます。
train_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, optional, defaults to (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。

For details on GKD, see the paper: On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.

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()を使用して再読み込みできます。

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

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 の元の引数を上書きするための書き込み権限を持つトークン。
revision (str, 省略可) — コミット元の git リビジョンです。“main” ブランチのヘッドがデフォルトです。
kwargs (dict[str, Any]、省略可) — ~Trainer.create_model_card に渡される追加のキーワード引数。

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

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) — 学生データ比率（つまり、オンポリシーの学生生成出力の割合）を制御する Lambda パラメータです。
beta (float, 任意、デフォルトは 0.5) — Generalized Jensen-Shannon Divergence ロスの 0.0 と 1.0 の間の補間係数です。 0.0 の場合、ロスは KL 分散になります。1.0 の場合、ロスは Inverse KL Divergence になります。
max_new_tokens (int、任意、デフォルトは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) — Sequence-Level KD を実行するかどうかを制御する Seq_kd パラメータです（教師モデルが生成した出力に対する教師あり 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