Personnalisation de l'entraînement

TRL est conçu avec une approche modulaire afin que les utilisateurs puissent personnaliser efficacement la boucle d'entraînement selon leurs besoins. Voici des exemples montrant comment appliquer et tester différentes techniques.

Bien que ces exemples utilisent le DPOTrainer, ces méthodes de personnalisation s’appliquent à la plupart (sinon à tous) les formateurs de TRL.

Utiliser différents optimiseurs et planificateurs

Par défaut, le DPOTrainer crée un optimiseur torch.optim.AdamW. Vous pouvez créer et définir un optimiseur différent et le passer à DPOTrainer comme suit :

from datasets import load_dataset
from torch import optim
from transformers import AutoModelForCausalLM
from trl import DPOTrainer

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
optimizer = optim.SGD(model.parameters(), lr=1e-6)

trainer = DPOTrainer(
    model=model,
    train_dataset=dataset,
    optimizers=(optimizer, None),
)
trainer.train()

Ajouter un planificateur de taux d'apprentissage

Vous pouvez également ajouter des planificateurs de taux d'apprentissage en transmettant à la fois l'optimiseur et le planificateur :

from torch import optim

optimizer = optim.AdamW(model.parameters(), lr=1e-6)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

trainer = DPOTrainer(..., optimizers=(optimizer, lr_scheduler))

Modèles de référence 8 bits

Puisque trl prend en charge tous les arguments par mot-clé lors du chargement d'un modèle depuis transformers à l'aide de from_pretrained, vous pouvez également tirer parti de load_in_8bit de transformers pour un fine-tuning plus efficace en mémoire.

En savoir plus sur le chargement de modèles en 8 bits dans transformers Charger en 8 bits ou 4 bits.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config=quantization_config)

trainer = DPOTrainer(..., ref_model=ref_model)

Ajouter des callbacks personnalisés

Vous pouvez personnaliser la boucle d’entraînement en ajoutant des callbacks pour la journalisation, la surveillance ou l’arrêt anticipé. Les callbacks vous permettent d’exécuter du code personnalisé à des points spécifiques pendant l’entraînement.

from transformers import TrainerCallback


class CustomLoggingCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None:
            print(f"Step {state.global_step}: {logs}")


trainer = DPOTrainer(..., callbacks=[CustomLoggingCallback()])

Ajouter des métriques d'évaluation personnalisées

Vous pouvez définir des métriques d'évaluation personnalisées à suivre pendant l'entraînement. Cela est utile pour surveiller les performances du modèle sur des tâches spécifiques.

def compute_metrics(eval_preds):
    logits, labels = eval_preds
    # Add your metric computation here
    return {"custom_metric": 0.0}


training_args = DPOConfig(..., eval_strategy="steps", eval_steps=100)

trainer = DPOTrainer(..., eval_dataset=eval_dataset, compute_metrics=compute_metrics)

Utiliser l’entraînement en précision mixte

L’entraînement en précision mixte peut accélérer considérablement l’entraînement et réduire l’utilisation de la mémoire. Vous pouvez l’activer en définissant bf16=True ou fp16=True dans la configuration d’entraînement.

# Use bfloat16 precision (recommended for modern GPUs)
training_args = DPOConfig(..., bf16=True)

Note : utilisez bf16=True pour les GPU Ampere (A100, RTX 30xx) ou plus récents, et fp16=True pour les GPU plus anciens.

Utiliser l'accumulation de gradients

Lors de l’entraînement avec une mémoire GPU limitée, l’accumulation de gradients permet de simuler des tailles de lot plus importantes en accumulant les gradients sur plusieurs étapes avant de mettre à jour les poids.

# Simulate a batch size of 32 with per_device_train_batch_size=4 and gradient_accumulation_steps=8
training_args = DPOConfig(
    ...,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

Changer l'objectif d'entraînement

Sous-classer un formateur est la méthode la plus simple pour personnaliser l’entraînement. La perte est définie dans la méthode compute_loss ; pour entraîner avec un objectif différent, il suffit de sous-classer le formateur et de redéfinir cette méthode. La préparation des données, la génération, la journalisation et la sauvegarde des points de contrôle sont hérités, de sorte que la sous-classe ne contient que ce qui change réellement.

import torch

from trl import DPOTrainer
from trl.trainer.utils import selective_log_softmax


class MyDPOTrainer(DPOTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # In this example: the hinge loss from https://huggingface.co/papers/2309.06657
        input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
        shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
        logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
        with torch.no_grad():
            ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits

        # Sum the log-probs over the completion tokens, for the policy and for the reference model
        logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
        ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)

        chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
        ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
        delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
        return torch.relu(1 - self.beta * delta).mean()

Ajouter des paramètres à la configuration

Sous-classez la config pour déclarer les paramètres dont votre perte a besoin :

from dataclasses import dataclass, field

from trl import DPOConfig


@dataclass
class MyDPOConfig(DPOConfig):
    my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})

Changer le format du lot

Lorsque la perte nécessite des entrées que le collateur par défaut ne produit pas, passez votre propre collateur et surchargez _prepare_dataset pour laisser le jeu de données intact s’il est déjà au format attendu :

class MyDPOTrainer(DPOTrainer):
    def _prepare_dataset(self, dataset, *args, **kwargs):
        return dataset


trainer = MyDPOTrainer(..., data_collator=my_collator)

Exemple complet

L’exemple de SFT par diffusion par blocs examples/sft_diffusion_gemma/sft_diffusion_gemma.py combine les trois : DiffusionGemmaSFTConfig étend SFTConfig avec les paramètres de canevas et de corruption, et DiffusionGemmaSFTTrainer étend SFTTrainer et remplace l’entropie croisée autorégressive par un objectif de débruitage par diffusion par blocs. Aucun des deux ne nécessite de modification de la bibliothèque.

Antidoom, un outil open source de Liquid AI, étend DPOTrainer avec l’optimisation des préférences sur le dernier jeton (Antislop), une perte de préférence sur une seule position de jeton qui réduit les boucles de répétition dans les modèles de raisonnement.