Subclassing de métodos do Trainer

Subclasses dos métodos de Trainer para alterar o comportamento do treinamento sem reescrever o loop inteiro. O subclassing modifica o training loop, por exemplo, a passagem para frente ou o cálculo da perda.

Antes de criar uma subclasse, considere se você precisa alterar o que o Trainer calcula ou quando e se ele age. Para lógica de tempo e condicional, use um Callback em vez disso. Callbacks controlam quando as coisas acontecem (log, avaliação, parada antecipada) e a criação de subclasse altera o que acontece (cálculo da perda, carregamento de dados, otimização).

Consulte a documentação da API do Trainer para obter a lista completa de métodos que você pode subclassear. Métodos privados (prefixados com _), como _save_checkpoint ou _evaluate, também podem ser sobrescritos, mas eles podem ser alterados sem aviso prévio.

get_train_dataloader

O método padrão get_train_dataloader() carrega um lote, treina com ele, descarta-o e carrega o próximo lote.

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

GRPO é um algoritmo de aprendizado por reforço online que gera complementos antes de treinar com eles. Gerar complementos a cada etapa é caro, pois é autoregressivo. Um complemento de 512 tokens requer ~512 passes de forward sequenciais, em comparação com um único pass de forward para uma etapa de treinamento. GRPOTrainer é uma subclasse de get_train_dataloader() para agrupar a geração em várias etapas.

trl.GRPOTrainer.get_train_dataloader carrega batches de prompts de geração para múltiplas etapas de treinamento de uma só vez, multiplicando o tamanho do lote pelo argumento steps_per_generation. Se train_batch_size=4 e steps_per_generation=8, o dataloader produz lotes de 32, reduzindo o custo de geração em 8x.

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

Calcular a perda

compute_loss() retorna a perda de entropia cruzada calculada pelo modelo.

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 mede o quão fortemente o modelo de política prefere uma resposta escolhida em relação a uma rejeitada, em comparação com um modelo de referência. DPOTrainer é uma subclasse de compute_loss() porque o cálculo da perda difere da entropia cruzada padrão em vários aspectos:

• o modelo nunca vê os rótulos; ele apenas retorna logits para que o DPO calcule as log-probabilidades a partir deles
• as respostas escolhidas e rejeitadas são concatenadas
• um modelo de referência calcula seus próprios log-probs
• a perda é uma função de π_chosen, π_rejected, π_ref_chosen, π_ref_rejected

Nenhum dos acima se encaixa no método padrão 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

Próximas etapas

• Para mais exemplos do mundo real, veja como GRPOTrainer e DPOTrainer estendem Trainer no TRL, ou como Axolotl cria treinadores personalizados com base nele.
• Consulte o guia de Callbacks se você precisar apenas personalizar o que acontece durante um evento de treinamento, como registrar métricas ao final de uma etapa de treinamento.