Fitur Trainer

Setiap resep di bawah ini mendemonstrasikan fitur Trainer tertentu: fungsi loss kustom, evaluasi yang hemat memori, strategi checkpointing, dan lainnya.

Buka issue jika ada fitur atau alur kerja yang ingin Anda lihat di sini.

Fungsi loss kustom

Lewatkan compute_loss_func ke Trainer untuk mengganti fungsi loss bawaan. Fungsi ini dijalankan setelah forward pass dan hanya mendefinisikan cara loss dihitung dari output. Untuk memodifikasi forward pass itu sendiri, subclass compute_loss() sebagai gantinya.

Fungsi loss kustom harus memiliki tanda tangan berikut:

import torch.nn.functional as F

def my_loss_fn(outputs, labels, num_items_in_batch):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

• outputs adalah output mentah dari model (outputs.logits memiliki bentuk (batch, seq_len, vocab_size)).
• labels adalah token id yang diambil dari batch input oleh Trainer sebelum forward pass.
• num_items_in_batch adalah jumlah target prediksi di seluruh batch yang terakumulasi. Untuk model LM kausal, ini menghitung label yang digeser (labels[..., 1:]), karena pergeseran label membuat posisi 0 dari setiap urutan tidak memiliki target. Lihat Loss scaling untuk detailnya. Trainer melewati normalisasi loss otomatis ketika fungsi loss kustom disediakan, sehingga fungsi Anda harus menangani normalisasi secara langsung.

trainer = Trainer(
    model=model,
    args=TrainingArguments(...),
    train_dataset=train_dataset,
    compute_loss_func=my_loss_fn,
)
trainer.train()

Lihat panduan subclassing untuk contoh lain tentang penggantian compute_loss().

Evaluasi saat mulai

Set eval_on_start=True untuk menjalankan satu putaran eval penuh sebelum langkah pelatihan pertama. Eval pra-pelatihan mengungkap masalah pada pipeline evaluasi lebih awal, terutama selama eksekusi yang panjang.

eval_on_start memerlukan eval_strategy yang valid (misalnya "epoch") dan dataset evaluasi.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_on_start=True,
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

Evaluasi penuh menambah waktu, sehingga paling berguna pada eksekusi pertama atau setelah memodifikasi compute_metrics.

Evaluasi hemat memori

Selama evaluasi, Trainer menjalankan forward pass pada setiap batch dan menggabungkan logits menjadi satu tensor di GPU. Setelah dataset eval diproses sepenuhnya, Trainer memindahkan logits yang telah digabungkan ke CPU dan memanggil compute_metrics. Untuk model besar atau set eval yang besar, logits yang terakumulasi dapat menghabiskan memori GPU meskipun pelatihan pada perangkat keras yang sama berjalan lancar, karena pelatihan hanya menyimpan satu batch aktivasi pada satu waktu.

eval_accumulation_steps

Offload prediksi yang terakumulasi dari GPU ke CPU setiap n batch. Nilai yang lebih rendah mengurangi memori GPU dengan biaya transfer CPU yang lebih sering.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,   # move predictions to CPU every 16 batches
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

preprocess_logits_for_metrics

Dipanggil sekali per batch eval di GPU, tepat setelah forward pass dan sebelum akumulasi logit. Nilai yang dikembalikan menggantikan logits di eval_pred.predictions. Menjalankan komputasi pada level batch mengurangi ukuran tensor per batch dan memberikan eval_accumulation_steps tensor yang lebih kecil untuk di-offload.

import evaluate
from transformers import Trainer, TrainingArguments

metric = evaluate.load("accuracy")

def preprocess_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple):
        logits = logits[0]
    return logits.argmax(dim=-1)

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    labels = labels[:, 1:].reshape(-1)
    preds = preds[:, :-1].reshape(-1)
    return metric.compute(predictions=preds, references=labels)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)
trainer.train()

Kinerja Dataloader

Secara default, Trainer membuat dataloader dengan dataloader_num_workers=0. Data dimuat di proses utama sementara GPU menganggur, yang terlihat sebagai utilisasi GPU yang rendah di antara batch.

Kedua dataloader_persistent_workers dan dataloader_prefetch_factor membutuhkan dataloader_num_workers > 0.

• dataloader_persistent_workers menjaga proses anak worker tetap aktif antar epoch untuk menghindari inisialisasi ulang dari awal, dengan konsekuensi penggunaan memori yang lebih tinggi.
• dataloader_prefetch_factor mengontrol berapa banyak batch yang disiapkan di muka oleh setiap worker. Dengan dataloader_prefetch_factor=2 dan num_workers=4, hingga 8 batch akan tersimpan di memori sementara GPU melatih pada batch saat ini.

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="out",
    dataloader_num_workers=4,            # spawn 4 worker subprocesses
    dataloader_persistent_workers=True,  # keep them alive between epochs
    dataloader_prefetch_factor=2,        # each worker preloads 2 batches ahead
)

Kelompokkan sampel berdasarkan panjang

Gunakan train_sampling_strategy="group_by_length" untuk mengelompokkan contoh dengan panjang yang serupa dan mengurangi padding. Ketika Anda tidak menyediakan panjang yang sudah dihitung sebelumnya, Trainer akan menyimpulkannya dari input model pertama di setiap item dataset. Hal ini juga berlaku ketika dataset multimodal berbasis processor mengembalikan objek BatchFeature, karena objek tersebut merupakan kontainer fitur yang mirip dengan mapping.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="qwen3-vl-finetuned",
    train_sampling_strategy="group_by_length",
)

Jika Dataset sudah memiliki kolom panjang yang telah dihitung sebelumnya, Trainer akan menggunakan kolom tersebut. Nama kolom default adalah length. Atur length_column_name jika dataset Anda menggunakan nama lain. Strategi ini memerlukan dataset dengan panjang yang diketahui dan diabaikan untuk IterableDataset.

Pengambilan sampel penyesuaian ulang batch

Pada dataset dengan panjang variabel, ketidakseimbangan di dalam micro-batch dan antar perangkat menyebabkan perangkat membuang waktu untuk padding dan menganggur pada langkah sinkronisasi gradien.

Setel train_sampling_strategy="batch_rebalance" di TrainingArguments untuk mengurangi kedua efek tersebut. Untuk setiap langkah optimizer, sampler:

1. Mengurutkan sampel dalam batch berdasarkan panjangnya.
2. Membagi batch yang telah diurutkan ke berbagai perangkat sehingga biaya token bantalan (padded-token) dari setiap micro-batch seimbang. Micro-batch dengan sampel panjang akan menerima lebih sedikit sampel, sedangkan micro-batch dengan sampel pendek akan menerima lebih banyak sampel.

Hal ini mengurangi padding di dalam setiap micro-batch, dan setiap perangkat menyelesaikan micro-batch pada waktu yang kurang lebih sama, sehingga mengurangi waktu idle saat sinkronisasi dan menjaga memori puncak tetap lebih rendah daripada "group_by_length". Strategi ini untuk saat ini hanya didukung untuk pelatihan data-parallel (tensor parallelism belum didukung).

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        train_sampling_strategy="batch_rebalance",  # balance padding cost across devices
        per_device_train_batch_size=8,              # average samples per micro-batch
        length_column_name="length",                # optional: dataset column with precomputed lengths
    ),
    train_dataset=train_dataset,
)
trainer.train()

per_device_train_batch_size di sini merupakan nilai rata-rata, bukan jumlah pasti per langkah: beberapa micro-batch memiliki lebih sedikit sampel dan beberapa memiliki lebih banyak, tetapi total jumlah sampel yang dilatih per langkah tetap sama seperti pada strategi sampling terdistribusi biasa.

Sampler membutuhkan panjang setiap sampel untuk mengurutkan dan menyeimbangkan batch. Secara default, Trainer memindai seluruh dataset sekali di awal pelatihan untuk menghitung panjang-panjang tersebut. Untuk melewati pemindaian ini, hitung terlebih dahulu panjang-panjangnya ke dalam kolom dataset (misalnya, selama pra-pemrosesan) dan lewati namanya sebagai length_column_name (default "length").

NEFTune

NEFTune menambahkan noise acak ke token embeddings selama forward pass. Noise ini berfungsi sebagai regularisasi dan dapat meningkatkan performa untuk instruction fine-tuning.

Aktifkan NEFTune dengan mengatur neftune_noise_alpha di TrainingArguments. Nilai alpha yang umum berkisar antara 5 hingga 15.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        num_train_epochs=3,
        neftune_noise_alpha=5,
    ),
    train_dataset=train_dataset,
)
trainer.train()

NEFTune hanya memengaruhi proses pelatihan, dan lapisan embedding asli dipulihkan setelah pelatihan selesai.

Pencatatan

Kendalikan kapan dan di mana Trainer menulis entri log dengan logging_strategy, logging_steps, dan report_to.

• logging_strategy="steps" mencatat setiap logging_steps() pembaruan optimizer. Gunakan "epoch" untuk mencatat di akhir setiap epoch.
• report_to mengalirkan log ke pelacak eksperimen seperti Trackio.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        logging_strategy="steps",
        logging_steps=50,               # write a log entry every 50 optimizer updates
        report_to="trackio",            # stream to Trackio (or "wandb", "tensorboard", …)
        run_name="model-experiment-v1", # display name in the tracker
    ),
    train_dataset=train_dataset,
)
trainer.train()

Checkpointing

Trainer menyimpan checkpoint setiap save_steps() pembaruan optimizer dan menyimpan semuanya (atau ~TrainingArguments.save_total_limit terbaru).

save_strategy="best" hanya menyimpan satu checkpoint terbaik berdasarkan metrik tertentu. Checkpoint baru disimpan hanya ketika metrik yang dipantau membaik, sehingga menghemat ruang disk dan mencegah penumpukan checkpoint yang sudah tidak relevan.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        save_strategy="best",
        metric_for_best_model="perplexity",   # save when eval perplexity improves
        greater_is_better=False,              # lower perplexity is better
        load_best_model_at_end=True,          # load the best weights after training finishes
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,          # must return {"perplexity": ...}
)
trainer.train()

Lanjutkan pelatihan

Lewati resume_from_checkpoint=True ke train() jika pelatihan terputus dan Anda ingin melanjutkan tanpa kehilangan progres. Pelatihan akan dilanjutkan dari checkpoint terbaru di output_dir.

trainer.train(resume_from_checkpoint=True)

Tentukan jalur checkpoint untuk melanjutkan dari titik tertentu.

trainer.train(resume_from_checkpoint="out/checkpoint-1000")

Saat melanjutkan, Trainer memulihkan state optimizer, state scheduler, dan state RNG.

Pemulihan checkpoint memerlukan file state optimizer dan scheduler di direktori checkpoint. Jika file-file tersebut tidak ada (misalnya, ketika save_only_model=True), optimizer akan dimulai ulang dari awal.

Checkpointing JIT

Dengan checkpointing berkala (save_strategy="steps" atau "epoch"), Anda akan kehilangan kemajuan pelatihan apa pun yang terjadi antara checkpoint terakhir yang disimpan dan saat interupsi. Pada klaster bersama dengan beban kerja yang dapat dipreempt seperti Kueue, pekerjaan dapat dihentikan kapan saja, sehingga selang waktu tersebut dapat berarti berjam-jam komputasi yang terbuang.

Checkpointing JIT (Just-In-Time) menutup celah ini. Ketika trainer menerima sinyal SIGTERM, ia menyimpan checkpoint pada titik persis saat pelatihan terinterupsi, sehingga Anda dapat melanjutkan dengan kehilangan progres yang minimal. Fitur ini bekerja berdampingan dengan checkpointing periodik. Penyimpanan periodik melindungi dari crash dan kegagalan perangkat keras, sementara penyimpanan JIT melindungi dari preemption dan shutdown yang anggun.

Aktifkan dengan mengatur enable_jit_checkpoint=True di TrainingArguments.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="your-model",
    enable_jit_checkpoint=True,
)

Ketika SIGTERM diterima, Trainer menunggu langkah pelatihan saat ini selesai, menyimpan checkpoint, dan menghentikan pelatihan secara mulus. File sentinel (checkpoint-is-incomplete.txt) ditulis saat penyimpanan dimulai dan dihapus setelah checkpoint sepenuhnya ditulis. Jika direktori checkpoint masih berisi file ini, penyimpanan terputus sebelum selesai. Trainer tidak memeriksanya secara otomatis, jadi periksa sendiri sebelum melanjutkan.

Lanjutkan dari checkpoint JIT dengan cara yang sama seperti checkpoint lainnya.

trainer.train(resume_from_checkpoint=True)

Anda harus mengonfigurasi orkestrator Anda untuk memberikan waktu yang cukup agar checkpoint dapat selesai. Periode graceful shutdown default Kubernetes hanya 30 detik, yang biasanya tidak cukup untuk model yang lebih besar.

spec:
  template:
    spec:
      terminationGracePeriodSeconds: 300

#SBATCH --signal=TERM@300

Hitung periode toleransi yang diperlukan sebagai waktu langkah pelatihan terpanjang yang mungkin ditambah waktu penyimpanan checkpoint, ditambah jeda kill_wait selama 3 detik sebelum checkpoint dimulai. Sebagai contoh, jika satu langkah pelatihan memakan waktu hingga 2 menit dan penyimpanan checkpoint memakan waktu 2 menit, tetapkan waktu toleransi minimal 243 detik.