Sous-classer les méthodes de Trainer

Sous-classez les méthodes de Trainer pour modifier le comportement de l’entraînement sans réécrire la boucle entière. Le sous-classage modifie la boucle d’entraînement, par exemple le passage avant ou le calcul de la perte.

Avant de créer une sous-classe, réfléchissez à savoir si vous devez modifier ce que Trainer calcule ou quand et si il agit. Pour la gestion du temps et la logique conditionnelle, utilisez plutôt un Callback. Les Callbacks contrôlent le moment où les événements se produisent (journalisation, évaluation, arrêt anticipé), tandis que la création d'une sous-classe modifie la nature de ces événements (calcul de la perte, chargement des données, optimisation).

Consultez la documentation de l’API Trainer pour obtenir la liste complète des méthodes que vous pouvez surcharger. Les méthodes privées (préfixées par _), comme _save_checkpoint ou _evaluate, peuvent également être surchargées, mais elles sont susceptibles d’évoluer sans préavis.

get_train_dataloader

La méthode standard get_train_dataloader() charge un lot, entraîne le modèle dessus, le supprime, puis charge le lot suivant.

def get_train_dataloader(self):
    return self._get_dataloader(
        batch_size=self._train_batch_size,
        ...
)

GRPO est un algorithme d'apprentissage par renforcement en ligne qui génère des complétions avant de s'entraîner dessus. Générer des complétions à chaque étape est coûteux car le processus est autorégressif. Une complétion de 512 jetons nécessite environ 512 passes avant séquentielles, contre une seule passe avant pour une étape d'entraînement. GRPOTrainer sous-classe get_train_dataloader() pour regrouper la génération sur plusieurs étapes.

trl.GRPOTrainer.get_train_dataloader charge des batches de prompts de génération pour plusieurs étapes d’entraînement à la fois en multipliant la taille du batch par l’argument steps_per_generation. Si train_batch_size=4 et steps_per_generation=8, le dataloader produit des batches de 32, réduisant le coût de génération par un facteur de 8.

def get_train_dataloader(self):
    dataloader_params = {
        "batch_size": self._train_batch_size * self.args.steps_per_generation, # this is the only change
        ...
    }

Calculer la perte

compute_loss() renvoie la perte d'entropie croisée calculée par le modèle.

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    ...
    outputs = model(**inputs)
    ...
    loss = outputs["loss"] # get loss from model

    return (loss, outputs) if return_outputs else loss

DPO mesure dans quelle mesure le modèle de politique préfère une réponse choisie à une réponse rejetée, par rapport à un modèle de référence. DPOTrainer est une sous-classe de compute_loss() car le calcul de la perte diffère de l'entropie croisée standard à plusieurs égards :

• le modèle ne voit jamais les étiquettes ; il ne renvoie que des logits pour que DPO puisse en calculer les log-probabilités
• les réponses choisies et rejetées sont concaténées
• un modèle de référence calcule ses propres log-probs
• la perte est une fonction de π_chosen, π_rejected, π_ref_chosen, π_ref_rejected

Aucune des méthodes ci-dessus ne correspond à la méthode standard Trainer.compute_loss().

def compute_loss(
    self,
    model: PreTrainedModel | nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs=False,
    num_items_in_batch=None,
) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]:
    ...
    outputs = model(**inputs)
    logits = outputs.logits
    logps = get_logps(logits, inputs)
    chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    ref_logits = self.ref_model(**inputs).logits
    ref_logps = get_logps(ref_logits, inputs)
    ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    chosen_scores = chosen_logps - ref_chosen_logps
    rejected_scores = rejected_logps - ref_rejected_logps
    per_sequence_loss = -F.logsigmoid(self.beta * chosen_scores - rejected_scores)
    loss = per_sequence_loss.mean()
    return (loss, outputs) if return_outputs else loss

Étapes suivantes

• Pour plus d’exemples concrets, consultez la manière dont GRPOTrainer et DPOTrainer étendent Trainer dans TRL, ou comment Axolotl construit des entraîneurs personnalisés à partir de celui-ci.
• Consultez le guide Callbacks si vous avez simplement besoin de personnaliser ce qui se produit lors d'un événement d'entraînement, comme l'enregistrement des métriques à la fin d'une étape d'entraînement.