Formateur de distillation de connaissances généralisée

Aperçu

La distillation de connaissances généralisée (GKD) a été proposée dans On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes par Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist et Olivier Bachem.

L'abstract de l'article est le suivant :

La distillation de connaissances (KD) est largement utilisée pour compresser un modèle enseignant afin de réduire son coût d'inférence et son empreinte mémoire, en entraînant un modèle étudiant plus petit. Cependant, les méthodes actuelles de KD pour les modèles de séquences auto-régressifs souffrent d'un décalage de distribution entre les séquences de sortie observées pendant l'entraînement et celles générées par l'étudiant lors de l'inférence. Pour remédier à ce problème, nous introduisons la distillation de connaissances généralisée (GKD). Au lieu de s'appuyer uniquement sur un ensemble fixe de séquences de sortie, GKD entraîne l'étudiant sur ses propres séquences de sortie auto-générées en exploitant les retours de l'enseignant sur ces séquences. Contrairement aux approches de KD supervisées, GKD offre également la flexibilité d'employer des fonctions de perte alternatives entre l'étudiant et l'enseignant, ce qui peut s'avérer utile lorsque l'étudiant manque d'expressivité pour imiter la distribution de l'enseignant. De plus, GKD facilite l'intégration transparente de la distillation avec l'ajustement fin par apprentissage par renforcement (RLHF). Nous démontrons l'efficacité de GKD pour la distillation de modèles de langage auto-régressifs sur des tâches de résumés, de traduction et de raisonnement arithmétique, ainsi que pour la distillation indépendante de la tâche pour l'ajustement fin par instructions.

Les aspects clés de GKD sont les suivants :

Elle traite le décalage de distribution entre l’entraînement et l’inférence dans les modèles de séquences auto-régressifs en entraînant le modèle élève sur les séquences de sortie qu’il génère lui-même.
GKD permet une flexibilité dans le choix de différentes mesures de divergence entre les modèles élève et enseignant via la divergence de Jensen-Shannon généralisée (JSD), ce qui peut s’avérer utile lorsque l’élève ne dispose pas de la capacité nécessaire pour imiter entièrement l’enseignant.

Cette méthode de post-entraînement a été proposée par Kashif Rasul et Lewis Tunstall.

Conseils d'utilisation

Le experimental.gkd.GKDTrainer est un wrapper autour de la classe SFTTrainer qui prend en argument un modèle enseignant. Il nécessite que trois paramètres soient définis via le experimental.gkd.GKDConfig, à savoir :

lmbda : contrôle la fraction de données étudiantes, c’est-à-dire la proportion de sorties générées par l’étudiant sur politique. Lorsque lmbda=0.0, la perte se réduit à un JSD supervisé où l’étudiant est entraîné avec les probabilités au niveau des jetons du professeur. Lorsque lmbda=1.0, la perte se réduit à un JSD sur politique, où l’étudiant génère des séquences de sortie et reçoit un retour spécifique aux jetons sur ces séquences de la part du professeur. Pour les valeurs comprises entre [0, 1], le choix est aléatoire entre les deux, basé sur la valeur de lmbda pour chaque lot.
seq_kd : contrôle l’exécution de la KD au niveau de la séquence (peut être considérée comme un FT supervisé sur les sorties générées par l’enseignant). Lorsque seq_kd=True et lmbda=0.0, la perte se réduit à une JSD supervisée, où l’enseignant génère des séquences de sortie et l’élève reçoit un retour spécifique à chaque jeton sur ces séquences de la part de l’enseignant.
beta : contrôle l'interpolation dans la divergence généralisée de Jensen-Shannon. Lorsque beta=0.0, la perte approxime la divergence KL avant, tandis que pour beta=1.0, elle approxime la divergence KL inverse. Pour les valeurs comprises entre [0, 1], elle interpole entre les deux.

Les auteurs constatent que les données on-policy (valeur élevée de lmbda) offrent de meilleures performances et que la valeur optimale de beta varie selon la tâche et la méthode d'évaluation.

Assurez-vous que attn_implementation="kernels-community/flash-attn2" lors de l'entraînement des modèles Gemma. Sinon, vous rencontrerez des NaN dans les logits en raison de la technique de soft capping adoptée par cette architecture.

L'API de base est la suivante :

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

Type de jeu de données attendu

Le jeu de données doit être formaté sous forme d'une liste de « messages », où chaque message est une liste de dictionnaires avec les clés suivantes :

role : soit system, soit assistant ou user
content : le contenu du message

GKDTrainer

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

Paramètres

model (PreTrainedModel ou torch.nn.Module ou str, facultatif) — Modèle à entraîner, ou l'identifiant en chaîne du modèle à instancier à partir d'un modèle pré-entraîné.
teacher_model (PreTrainedModel ou torch.nn.Module ou str, facultatif) — Modèle enseignant pour la distillation de connaissances, ou l'identifiant en chaîne du modèle à instancier à partir d'un modèle pré-entraîné.
args (experimental.gkd.GKDConfig, facultatif) — Arguments d'entraînement.
data_collator (DataCollator, facultatif) — Collateur de données pour regrouper les échantillons du jeu de données en lots. Par défaut, il s’agit d’un experimental.utils.DataCollatorForChatML utilisant le processing_class.
train_dataset (Dataset, facultatif) — Jeu de données pour l'entraînement.
eval_dataset (Dataset ou dict de Dataset, facultatif) — Jeu de données pour l'évaluation.
processing_class (PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin ou ProcessorMixin, facultatif) — Classe pour traiter les données.
compute_metrics (Callable, facultatif) — Fonction permettant de calculer les métriques lors de l'évaluation. Elle doit prendre en entrée un EvalPrediction et renvoyer un dictionnaire de chaînes de caractères vers des flottants.
callbacks (list de TrainerCallback, facultatif) — Callbacks à utiliser pendant l'entraînement.
optimizers (tuple de torch.optim.Optimizer et torch.optim.lr_scheduler.LambdaLR, facultatif, par défaut (None, None)) — Tuple contenant l'optimiseur et le planificateur de taux d'apprentissage à utiliser pour l'entraînement.
preprocess_logits_for_metrics (Callable, facultatif) — Fonction permettant de prétraiter les logits avant le calcul des métriques. Elle doit prendre en entrée les logits et les labels et renvoyer les logits à utiliser pour le calcul des métriques.
peft_config (PeftConfig, facultatif) — Configuration PEFT à utiliser pour l'entraînement avec PEFT. Si None, PEFT n'est pas utilisé. S'il est fourni, le model sera enveloppé avec l'adaptateur PEFT spécifié.
formatting_func (Callable, facultatif) — Fonction pour formater le jeu de données. Doit prendre un exemple en entrée et renvoyer un exemple.

Trainer pour la distillation de connaissances généralisée (GKD) des modèles de langage.

Pour plus de détails sur GKD, consultez l'article : On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.

train

< source >

Paramètres

resume_from_checkpoint (str ou bool, facultatif) — Si c'est une str, chemin local d'un point de contrôle enregistré par une instance précédente de Trainer. Si c'est un bool et vaut True, charge le dernier point de contrôle dans args.output_dir enregistré par une instance précédente de Trainer. S'il est présent, l'entraînement reprendra à partir des états du modèle/optimiseur/planificateur chargés ici.
trial (optuna.Trial ou dict[str, Any], facultatif) — L’exécution de l’essai ou le dictionnaire d’hyperparamètres pour la recherche d’hyperparamètres.
ignore_keys_for_eval (list[str], facultatif) — Une liste de clés dans la sortie de votre modèle (s’il s’agit d’un dictionnaire) qui doivent être ignorées lors de la collecte des prédictions pour l’évaluation pendant l’entraînement.

Retourne

~trainer_utils.TrainOutput

Objet contenant le nombre d'étapes globales, la perte d'entraînement et les métriques.

Point d'entrée principal de l'entraînement.

save_model

< source >

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

Enregistrera le modèle, ce qui vous permettra de le recharger à l'aide de from_pretrained().

N'enregistrera que depuis le processus principal.

push_to_hub

< source >

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

Paramètres

commit_message (str, facultatif, par défaut "End of training") — Message à committer lors de la poussée.
blocking (bool, facultatif, par défaut True) — Indique si la fonction doit ne renvoyer que lorsque le git push est terminé.
token (str, facultatif, par défaut None) — Jeton avec permission d’écriture pour écraser les arguments originaux du Trainer.
revision (str, facultatif) — La révision git à partir de laquelle effectuer le commit. Par défaut, il s'agit de la tête de la branche « main ».
kwargs (dict[str, Any], facultatif) — Arguments supplémentaires passés à ~Trainer.create_model_card.

Téléchargez self.model et self.processing_class sur le hub de modèles 🤗 dans le dépôt self.args.hub_model_id.

GKDConfig

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

Paramètres

temperature (float, facultatif, par défaut 0.9) — Température pour l'échantillonnage. Plus la température est élevée, plus les complétions sont aléatoires.
lmbda (float, facultatif, valeur par défaut 0.5) — Paramètre Lambda qui contrôle la fraction de données étudiantes (c’est-à-dire la proportion de sorties générées par l’étudiant selon sa propre politique).
beta (float, facultatif, par défaut 0.5) — Coefficient d'interpolation entre 0.0 et 1.0 de la perte de divergence de Jensen-Shannon généralisée. Lorsque beta vaut 0.0, la perte est la divergence KL. Lorsque beta vaut 1.0, la perte est la divergence KL inverse.
max_new_tokens (int, facultatif, par défaut 128) — Nombre maximal de jetons à générer par complétion.
teacher_model_name_or_path (str, facultatif) — Nom ou chemin du modèle enseignant. Si None, le modèle enseignant sera identique au modèle en cours d’entraînement.
teacher_model_init_kwargs (dict[str, Any], facultatif) — Arguments par mot-clé à passer à AutoModelForCausalLM.from_pretrained lors de l'instantiation du modèle enseignant à partir d'une chaîne de caractères.
disable_dropout (bool, facultatif, valeur par défaut True) — Indique s'il faut désactiver le dropout dans le modèle.
seq_kd (bool, facultatif, valeur par défaut False) — Paramètre qui contrôle l’exécution du KD au niveau de la séquence (peut être considéré comme un FT supervisé sur la sortie générée par l’enseignant).

Classe de configuration pour experimental.gkd.GKDTrainer.

Cette classe ne contient que les paramètres spécifiques à l’entraînement GKD. Pour la liste complète des arguments d’entraînement, veuillez consulter la documentation de TrainingArguments et SFTConfig.

Update on GitHub

TRL

Formateur de distillation de connaissances généralisée

Aperçu

Conseils d'utilisation

Type de jeu de données attendu

GKDTrainer

classe trl.experimental.gkd.GKDTrainer

train

save_model

push_to_hub

GKDConfig

classe trl.experimental.gkd.GKDConfig