子类化 Trainer 方法

通过子类化 Trainer 的方法来更改训练行为，而无需重写整个循环。子类化会修改 training loop，例如前向传播或损失计算。

在创建子类之前，请先考虑你需要改变的是 Trainer 计算什么，还是它何时以及是否执行操作。对于时序和条件逻辑，请改用 Callback。Callback 控制事件发生的时机（如日志记录、评估、早停），而子类化则改变事件发生的内容（如损失计算、数据加载、优化）。

请参阅 Trainer API 文档，以获取可子类化的完整方法列表。私有方法（以 _ 为前缀）例如 _save_checkpoint 或 _evaluate 也可以被重写，但这些方法可能会在不另行通知的情况下发生变更。

get_train_dataloader

标准的 get_train_dataloader() 方法会加载一个批次，在其上进行训练，然后丢弃该批次，再加载下一个批次。

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

GRPO 是一种在线强化学习算法，它先生成补全，再基于这些补全进行训练。由于生成过程是自回归的，每一步都生成补全的开销很大：生成一个 512 个 token 的补全大约需要 512 次顺序前向传播，而一个训练步骤只需一次前向传播。GRPOTrainer 通过重写 get_train_dataloader()，将生成过程跨多个步骤进行批处理。

trl.GRPOTrainer.get_train_dataloader 通过将 batch size 乘以 steps_per_generation 参数，一次性加载用于多个训练步骤的 batches 生成提示。如果 train_batch_size=4 且 steps_per_generation=8，dataloader 将生成大小为 32 的批次，从而将生成成本降低 8 倍。

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

计算损失

compute_loss() 返回模型计算得出的交叉熵损失。

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 衡量策略模型相对于参考模型，对优选回复的偏好程度高于被拒回复的程度。DPOTrainer 重写了 compute_loss()，因为损失计算在多个方面与标准交叉熵不同：

模型从未看到标签；它仅返回 logits，供 DPO 从中计算对数概率
选中的回复和被拒绝的回复会被拼接在一起
参考模型计算其自身的对数概率
损失是 π_chosen、π_rejected、π_ref_chosen 和 π_ref_rejected 的函数

以上均不符合标准的 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

后续步骤

如需更多实际示例，请参阅 TRL 中 GRPOTrainer 和 DPOTrainer 如何扩展 Trainer，或 Axolotl 如何在其基础上构建自定义训练器。
如果您只需要自定义训练事件（例如在训练步骤结束时记录指标）期间发生的内容，请查阅 Callbacks 指南。

Update on GitHub

Transformers

子类化 Trainer 方法

get_train_dataloader

计算损失

后续步骤