Akselerasi gradien

Batch besar menghasilkan aktivasi besar yang menguras memori GPU. Gradient accumulation memungkinkan Anda melatih dengan ukuran batch efektif yang lebih besar dengan menyebarkan komputasi gradien ke beberapa mini-batch.

Gradien terakumulasi di sepanjang n mini-batch sebelum optimizer memperbarui bobot. Sebagai contoh, dengan ukuran batch per perangkat sebesar 8 dan 4 langkah akumulasi, ukuran batch efektifnya adalah 32.

Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()

Gunakan akumulasi gradien hanya ketika batch yang lebih besar tidak muat di memori. Metode ini tidak meningkatkan throughput dibandingkan pelatihan dengan batch besar yang sesungguhnya.

Akhumulasi gradien selama gradient_accumulation_steps di seluruh per_device_train_batch_size.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)

Skala kerugian

Untuk fungsi loss kustom, sertakan num_items_in_batch agar Trainer membagi loss dengan jumlah target prediksi di seluruh mini-batch. Ini menormalisasi berdasarkan token, bukan jumlah langkah tetap dengan gradient_accumulation_steps. Jika tidak, Trainer membagi loss dengan gradient_accumulation_steps.

import torch.nn.functional as F

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

Untuk model LM kausal, num_items_in_batch menghitung label yang di-geser. Loss menggeser label sehingga prediksi pada posisi i menargetkan token pada posisi i + 1, yang membuat posisi 0 dari setiap urutan tidak memiliki target. Trainer mengecualikan posisi-posisi tersebut dan menghitung berdasarkan labels[..., 1:], sehingga penyebutnya sesuai dengan jumlah target prediksi yang digunakan oleh loss. Ketika data collator menyediakan shift_labels secara langsung, seperti collator tanpa padding, Trainer menghitung berdasarkan tensor tersebut. Jenis loss lainnya, seperti masked LM dan klasifikasi, menghitung tensor label secara penuh.

Jika Anda membuat subclass dari Trainer dan menimpa compute_loss(), atur loss_is_scaled_for_ga untuk menyatakan bagaimana loss Anda dinormalisasi: False untuk rata-rata per-batch, yang kemudian dibagi oleh Trainer dengan gradient_accumulation_steps, atau True jika Anda sudah membaginya dengan num_items_in_batch.

class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch

Langkah selanjutnya

• Baca dokumen penggunaan memori GPU untuk memahami apa yang mendorong penggunaan memori pada GPU selama pelatihan.
• Lihat panduan Gradient checkpointing untuk mempelajari cara mengurangi memori aktivasi dengan menghitung ulang aktivasi alih-alih menyimpannya di cache.
• Lihat panduan Mixed precision training untuk mempelajari cara menggunakan tipe data presisi lebih rendah guna mengurangi memori dan mempercepat pelatihan.
• Baca artikel blog Gradient Accumulation Fix untuk mempelajari cara akumulasi gradien dihitung.