Kustomisasi pelatihan

TRL dirancang dengan modularitas sehingga pengguna dapat menyesuaikan loop pelatihan secara efisien sesuai kebutuhan mereka. Di bawah ini adalah contoh cara menerapkan dan menguji berbagai teknik.

Meskipun contoh-contoh ini menggunakan DPOTrainer, metode kustomisasi ini berlaku untuk sebagian besar (jika tidak semua) trainer di TRL.

Gunakan optimizer dan scheduler yang berbeda

Secara default, DPOTrainer membuat optimizer torch.optim.AdamW. Anda dapat membuat dan mendefinisikan optimizer yang berbeda lalu meneruskannya ke DPOTrainer sebagai berikut:

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()

Tambahkan penjadwal laju pembelajaran

Anda juga dapat menambahkan penjadwal laju pembelajaran dengan meneruskan optimizer dan scheduler:

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))

Model referensi 8-bit

Karena trl mendukung semua argumen kata kunci saat memuat model dari transformers menggunakan from_pretrained, Anda juga dapat memanfaatkan load_in_8bit dari transformers untuk fine-tuning yang lebih hemat memori.

Baca lebih lanjut tentang pemuatan model 8-bit di transformers Load in 8bit or 4bit.

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)

Tambahkan callback kustom

Anda dapat menyesuaikan loop pelatihan dengan menambahkan callback untuk logging, monitoring, atau early stopping. Callback memungkinkan Anda untuk menjalankan kode kustom pada titik-titik tertentu selama pelatihan.

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()])

Tambahkan metrik evaluasi kustom

Anda dapat mendefinisikan metrik evaluasi kustom untuk dilacak selama pelatihan. Hal ini berguna untuk memantau kinerja model pada tugas tertentu.

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)

Gunakan pelatihan presisi campuran

Pelatihan presisi campuran dapat mempercepat pelatihan secara signifikan dan mengurangi penggunaan memori. Anda dapat mengaktifkannya dengan mengatur bf16=True atau fp16=True dalam konfigurasi pelatihan.

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

Catatan: Gunakan bf16=True untuk GPU Ampere (A100, RTX 30xx) atau yang lebih baru, dan fp16=True untuk GPU yang lebih lama.

Gunakan akumulasi gradien

Saat melatih dengan memori GPU terbatas, akumulasi gradien memungkinkan Anda mensimulasikan ukuran batch yang lebih besar dengan mengakumulasi gradien selama beberapa langkah sebelum memperbarui bobot.

# 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,
)

Ubah tujuan pelatihan

Membuat subclass dari trainer adalah cara termudah untuk menyesuaikan pelatihan. Loss didefinisikan dalam metode compute_loss, sehingga untuk melatih dengan tujuan yang berbeda, buat subclass dari trainer dan override metode tersebut. Persiapan data, generasi, logging, dan checkpointing diwarisi, sehingga subclass hanya berisi bagian yang benar-benar berubah.

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()

Tambahkan parameter ke config

Subclass config untuk mendeklarasikan parameter yang dibutuhkan loss Anda:

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."})

Ubah format batch

Ketika loss membutuhkan input yang tidak dihasilkan oleh collator default, berikan collator Anda sendiri, dan override _prepare_dataset agar dataset tetap tidak diubah jika sudah dalam format yang diharapkan:

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


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

Contoh lengkap

Contoh SFT block-diffusion examples/sft_diffusion_gemma/sft_diffusion_gemma.py menggabungkan ketiganya: DiffusionGemmaSFTConfig memperluas SFTConfig dengan parameter canvas dan corruption, dan DiffusionGemmaSFTTrainer memperluas SFTTrainer serta mengganti cross-entropy autoregressive dengan tujuan denoising block-diffusion. Keduanya tidak memerlukan perubahan pada library.

Antidoom, sebuah alat open-source dari Liquid AI, memperluas DPOTrainer dengan Final Token Preference Optimization (Antislop), yaitu loss preferensi pada satu posisi token yang mengurangi loop pengulangan pada model penalaran.

Update on GitHub