Checkpointing gradien

Forward pass biasanya meng-cache semua aktivasi perantara untuk digunakan kembali pada backward pass. Namun, aktivasi berskala dengan ukuran batch dan panjang sekuens. Gradient checkpointing hanya menyimpan aktivasi tertentu dan membuang sisanya. Hal ini memaksa backward pass untuk menghitung ulang beberapa aktivasi secara on-the-fly saat dibutuhkan.

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

Pelatihan biasanya lebih lambat karena jalur mundur (backward pass) menghitung ulang aktivasi yang dibuang, tetapi checkpointing mengurangi memori aktivasi.

Setel gradient_checkpointing=True untuk mengaktifkannya.

Gunakan dengan akumulasi gradien untuk mengurangi penggunaan memori lebih lanjut.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

Checkpointing parsial

Full gradient checkpointing menghitung ulang setiap layer yang dapat di-checkpoint. Jika eksekusi Anda memiliki sisa memori, checkpoint lebih sedikit layer untuk menukar sebagian penghematan memori dengan kecepatan.

Lewatkan every_n_layers ke gradient_checkpointing_enable() untuk memilih interval checkpointing.

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

Dengan every_n_layers=2, lapisan pertama dan setiap lapisan kedua setelahnya akan di-checkpoint. Lapisan yang di-checkpoint membuang aktivasi mereka selama forward pass dan menghitungnya ulang selama backward pass, sementara lapisan lainnya menyimpan aktivasi mereka di memori.

model.gradient_checkpointing_enable(every_n_layers=2)

Nilai default, every_n_layers=1, membuat checkpoint untuk setiap layer. Nilai yang lebih besar membuat checkpoint untuk layer pertama dan kemudian setiap n layer setelahnya, sehingga aktivasi layer lainnya tetap berada di memori. Sebagai contoh, every_n_layers=2 membuat checkpoint untuk layer 1, 3, 5, dan seterusnya. Hanya modul yang mewarisi dari GradientCheckpointingLayer yang dihitung. Modul lain yang mendukung gradient checkpointing tetap aktif.

Untuk menggunakan partial gradient checkpointing dengan Trainer, atur every_n_layers di gradient_checkpointing_kwargs.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

Offloading aktivasi yang disimpan

Gradient checkpointing menyimpan satu aktivasi per lapisan yang di-checkpoint di GPU. Pada panjang sekuens yang panjang, hal ini dapat mengonsumsi memori yang signifikan (sekitar layers x sequence x hidden x bytes_per_element). Atur offload untuk menyimpan aktivasi tersebut di memori host yang di-pin, dengan biaya salinan dari perangkat ke host selama forward pass dan salinan dari host ke perangkat selama backward pass.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

Kedua salinan tersebut berjalan di stream komputasi, sehingga ini menukar langkah yang lebih lambat dengan memori. Gunakan ini ketika sebuah run tidak muat dengan cara lain, bukan untuk mempercepatnya.

Langkah selanjutnya

• Baca dokumen penggunaan memori GPU untuk memahami apa yang mendorong penggunaan memori pada GPU selama pelatihan.
• Lihat panduan Mixed precision training untuk mempelajari cara menggunakan tipe data presisi yang lebih rendah guna mengurangi memori lebih lanjut dan mempercepat pelatihan.
• Lihat panduan Kernels untuk mempelajari cara mempercepat pelatihan dengan kernel fused kustom.