Membuat Subkelas Metode Trainer

Subclass metode Trainer untuk mengubah perilaku pelatihan tanpa menulis ulang seluruh loop. Subclassing memodifikasi training loop, misalnya forward pass atau perhitungan loss.

Sebelum melakukan subclassing, pertimbangkan apakah Anda perlu mengubah apa yang dihitung oleh Trainer atau kapan dan apakah ia bertindak. Untuk logika penjadwalan dan kondisional, gunakan Callback sebagai gantinya. Callback mengontrol kapan sesuatu terjadi (logging, evaluasi, early stopping), sedangkan subclassing mengubah apa yang terjadi (perhitungan loss, pemuatan data, optimasi).

Lihat dokumentasi API Trainer untuk daftar lengkap metode yang dapat Anda subclass. Metode privat (diawali dengan _) seperti _save_checkpoint atau _evaluate juga dapat di-override, tetapi metode ini dapat berubah tanpa pemberitahuan.

get_train_dataloader

Metode standar get_train_dataloader() memuat satu batch, melatih model berdasarkan batch tersebut, membuangnya, lalu memuat batch berikutnya.

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

GRPO adalah algoritma reinforcement learning online yang menghasilkan kelengkapan (completions) sebelum dilatih di atasnya. Menghasilkan kelengkapan pada setiap langkah sangat mahal karena bersifat autoregressive. Kelengkapan 512-token membutuhkan sekitar 512 forward pass berurutan, dibandingkan dengan satu forward pass untuk satu langkah pelatihan. GRPOTrainer mensubkelas get_train_dataloader() untuk melakukan pengelompokan (batching) generasi di antara beberapa langkah.

trl.GRPOTrainer.get_train_dataloader memuat batches prompt generasi untuk beberapa langkah pelatihan sekaligus dengan mengalikan ukuran batch dengan argumen steps_per_generation. Jika train_batch_size=4 dan steps_per_generation=8, dataloader menghasilkan batch berukuran 32, sehingga mengurangi biaya generasi sebanyak 8 kali lipat.

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

Hitung loss

compute_loss() mengembalikan loss entropi silang yang dihitung oleh model.

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 mengukur seberapa kuat model kebijakan memilih respons yang dipilih dibandingkan respons yang ditolak, relatif terhadap model referensi. DPOTrainer mensubkelas compute_loss() karena perhitungan loss-nya berbeda dari cross-entropy standar dalam beberapa hal:

• model tidak pernah melihat label; ia hanya mengembalikan logits agar DPO dapat menghitung log-probs darinya
• respons yang dipilih dan ditolak digabungkan
• model referensi menghitung log-probs-nya sendiri
• loss merupakan fungsi dari π_chosen, π_rejected, π_ref_chosen, π_ref_rejected

Tidak ada dari yang di atas yang sesuai dengan metode standar 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

Langkah selanjutnya

• Untuk contoh dunia nyata lainnya, lihat bagaimana GRPOTrainer dan DPOTrainer memperluas Trainer di TRL, atau bagaimana Axolotl membangun trainer kustom di atasnya.
• Periksa panduan Callbacks jika Anda hanya perlu menyesuaikan apa yang terjadi selama peristiwa pelatihan, seperti pencatatan metrik di akhir langkah pelatihan.