Mengurangi Penggunaan Memori

Alur kerja pelatihan sering kali dapat dioptimalkan untuk mengurangi konsumsi memori, dan TRL menyediakan beberapa fitur bawaan untuk membantu mencapainya.

Di bawah ini, kami menguraikan teknik-teknik tersebut dan merekomendasikan untuk bereksperimen dengan berbagai kombinasi untuk menentukan konfigurasi mana yang paling sesuai untuk pengaturan spesifik Anda.

Setiap metode mencakup contoh untuk trainer yang didukung. Jika Anda tidak yakin apakah suatu teknik kompatibel dengan trainer Anda, silakan lihat dokumentasi trainer yang sesuai.

Untuk strategi tambahan, seperti gradient checkpointing, yang didukung oleh semua trainer, lihat panduan kinerja transformers.

Pemotongan

Panjang sekuens dalam dataset dapat bervariasi secara signifikan. Ketika data di-batch, sekuens di-pad agar sesuai dengan sekuens terpanjang dalam batch tersebut, yang dapat menyebabkan penggunaan memori tinggi, bahkan jika sebagian besar sekuens relatif pendek.

Untuk mengurangi penggunaan memori, penting untuk memangkas urutan hingga panjang yang wajar. Meskipun pelatih TRL memangkas urutan secara bawaan, Anda mungkin ingin menyesuaikan panjang pemangkasan bawaan agar lebih sesuai dengan kasus penggunaan spesifik Anda.

from trl import DPOConfig

training_args = DPOConfig(..., max_length=...)

from trl import SFTConfig

training_args = SFTConfig(..., max_length=...)

Bagaimana cara memilih nilai max_length?

Jika max_length terlalu kecil, sebagian besar token Anda akan dibuang dan tidak berkontribusi pada pelatihan. Jika terlalu besar, penggunaan memori dapat melonjak, yang berpotensi menyebabkan error out-of-memory (OOM). Tanpa packing atau padding-free, max_length yang besar juga dapat menghasilkan pelatihan yang tidak efisien, karena banyak token yang berupa padding.

Untuk membantu Anda memilih nilai yang tepat, kami menyediakan utilitas untuk memvisualisasikan distribusi panjang sekuens dalam dataset Anda.

Packing

Teknik ini hanya tersedia untuk pelatihan SFT dan konfigurasi yang menggunakan FlashAttention (atau variasinya).

Truncation memiliki beberapa kerugian:

1. Kehilangan informasi: Token penting di akhir urutan dapat dibuang.
2. Memilih panjang pemotongan: Terlalu pendek menyebabkan hilangnya data; terlalu panjang mengurangi efisiensi.

Packing memitigasi masalah-masalah ini dengan mengelompokkan beberapa urutan ke dalam baris pelatihan yang sama, mengisi setiap baris hingga max_length.

TRL menerapkan packing menggunakan bin packing Best-Fit Decreasing (BFD), yang mengelompokkan urutan secara efisien sambil meminimalkan padding. Ketika suatu urutan melebihi max_length, strategi yang berbeda menentukan bagaimana token yang meluap ditangani.

TRL mendukung tiga strategi:

• "bfd" (default): Menggunakan Best-Fit Decreasing packing. Jika sebuah urutan melebihi max_length, token yang berlebih akan dibuang.

• "bfd_split": Menggunakan Best-Fit Decreasing packing, tetapi urutan panjang dipecah menjadi chunk ≤ max_length sebelum packing. Ini mempertahankan semua token dan mengikuti pendekatan yang diusulkan dalam Fewer Truncations Improve Language Modeling.

• "wrapped": Semua token digabungkan menjadi satu aliran lalu dipecah menjadi blok dengan panjang tetap. Strategi ini meminimalkan padding tetapi dapat mencampur contoh yang tidak terkait. Strategi ini sesuai dengan pra-pemrosesan concatenate-then-split yang dijelaskan dalam literatur (misalnya, Fewer Truncations Improve Language Modeling). Kelemahannya adalah memutus kontinuitas urutan untuk sebagian besar dataset, yang merugikan performa, sebagaimana dibahas dalam Qwen3-Coder-Next Technical Report.

Jika semua sekuens lebih pendek dari max_length, bfd dan bfd_split berperilaku identik, karena tidak diperlukan pemotongan atau pemisahan.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

PEFT untuk penyesuaian halus yang efisien parameter

Metode Fine-Tuning yang Efisien Parameter (PEFT) seperti LoRA termasuk teknik paling efektif untuk mengurangi penggunaan memori selama pelatihan. Alih-alih melatih semua parameter model, metode PEFT hanya melatih sejumlah kecil parameter adapter, yang secara signifikan mengurangi kebutuhan memori dan memungkinkan fine-tuning model yang lebih besar pada perangkat keras terbatas.

Untuk detail komprehensif tentang penggunaan PEFT dengan TRL, termasuk berbagai metode adapter, opsi kuantisasi, dan konfigurasi lanjutan, lihat Integrasi PEFT.

Untuk menggunakan PEFT guna mengurangi penggunaan memori:

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

PEFT dapat dikombinasikan dengan teknik pengurangan memori lainnya seperti kuantisasi (4-bit atau 8-bit) untuk penghematan memori yang lebih besar. Lihat Integrasi PEFT untuk contoh kuantisasi.

Tanpa padding

Batching tanpa padding adalah pendekatan alternatif untuk mengurangi penggunaan memori. Dalam metode ini, sebuah batch pertama-tama diambil sampelnya dan kemudian diratakan menjadi satu urutan tunggal, sehingga menghindari padding. Berbeda dengan packing, yang dapat menghasilkan urutan tidak lengkap dengan menggabungkan bagian dari sampel yang berbeda, batching tanpa padding memastikan bahwa semua urutan tetap utuh dan lengkap.

from trl import DPOConfig

training_args = DPOConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

from trl import SFTConfig

training_args = SFTConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

Offloading aktivasi

Activation offloading adalah teknik efisiensi memori yang mengurangi penggunaan VRAM GPU dengan memindahkan tensor aktivasi ke RAM CPU secara sementara selama forward pass, lalu mengembalikannya hanya ketika diperlukan untuk backward pass. Hal ini secara signifikan mengurangi penggunaan memori puncak dengan biaya sedikitnya peningkatan waktu pelatihan.

Untuk mengaktifkan pengalihan aktivasi dalam konfigurasi pelatihan SFT Anda:

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

Di balik layar, activation offloading mengimplementasikan saved_tensors_hooks dari PyTorch untuk memotong aktivasi selama forward pass. Fitur ini secara cerdas mengelola tensor mana yang akan di-offload berdasarkan ukuran dan konteks, serta menghindari offloading tensor output yang tidak efisien. Untuk optimasi performa, fitur ini dapat, melalui sebuah flag (yang bernilai true secara default), menggunakan CUDA streams untuk menumpuk komputasi dengan transfer CPU-GPU.

Menambahkan Padding pada Urutan ke Kelipatan

Teknik ini saat ini didukung untuk trainer SFT dan Reward.

Ketika diaktifkan, opsi ini memastikan bahwa semua urutan diisi hingga kelipatan dari nilai yang ditentukan.
Hal ini dapat meningkatkan efisiensi komputasi pada beberapa perangkat keras dengan menyelaraskan panjang urutan ke batas memori yang lebih ramah.

from trl import SFTConfig

training_args = SFTConfig(..., pad_to_multiple_of=2048)

from trl import RewardConfig

training_args = RewardConfig(..., pad_to_multiple_of=2048)

Alokator cache PyTorch

Pada eksekusi panjang, terutama RL online (GRPO, RLOO, Online DPO), memori GPU dapat terfragmentasi. Pengaturan expandable_segments:True memungkinkan caching allocator PyTorch untuk memperbesar segmen yang sudah ada, yang mengurangi selisih antara memori yang dialokasikan dan memori yang dicadangkan:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

Dengan mode tidur vLLM (vllm_enable_sleep_mode=True), vLLM hanya membaca PYTORCH_CUDA_ALLOC_CONF untuk mengakali pengaturan ini, jadi jangan hanya mengatur PYTORCH_ALLOC_CONF.

Menonaktifkan pengumpulan model untuk generasi dalam metode online

Saat menggunakan DeepSpeed ZeRO-3, bobot model dipecah (sharded) di beberapa GPU. Metode online melibatkan pembuatan kelengkapan (completions) dari model sebagai bagian dari proses pelatihan. Selama langkah ini, bobot model dikumpulkan secara sementara di satu GPU untuk generasi. Untuk model yang sangat besar, pengumpulan ini dapat menyebabkan error OOM, seperti yang dijelaskan dalam isu ini: #2250.

Jika Anda menemui masalah ini, Anda dapat menonaktifkan pengumpulan bobot model untuk generasi dengan mengatur parameter berikut:

from trl import GRPOConfig

training_args = GRPOConfig(..., ds3_gather_for_generation=False)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(..., ds3_gather_for_generation=False)

from trl import RLOOConfig

training_args = RLOOConfig(..., ds3_gather_for_generation=False)

Penyesuaian ini mencegah pengumpulan bobot model, sehingga menghindari error OOM, tetapi dapat menyebabkan kecepatan generasi yang lebih lambat.

Mode tidur vLLM

Saat menggunakan vLLM sebagai backend generasi untuk metode pelatihan online, Anda dapat mengaktifkan sleep mode untuk memindahkan parameter dan cache vLLM ke RAM CPU selama langkah optimasi, lalu memuatnya kembali ke VRAM GPU saat diperlukan untuk sinkronisasi bobot dan generasi.

from trl import GRPOConfig

training_args = GRPOConfig(..., vllm_enable_sleep_mode=True)

from trl import RLOOConfig

training_args = RLOOConfig(..., vllm_enable_sleep_mode=True)

Offloading bobot dan cache vLLM membantu menjaga penggunaan memori GPU tetap rendah, yang dapat sangat bermanfaat saat melatih model besar atau menggunakan sumber daya GPU yang terbatas. Namun, membangunkan mesin vLLM dari mode tidur memperkenalkan beberapa latensi transfer host–perangkat, yang dapat sedikit memengaruhi kecepatan pelatihan.

Checkpointing gradien

Gradient checkpointing menukar komputasi dengan memori dengan tidak menyimpan semua aktivasi perantara selama forward pass, dan menghitungnya ulang selama backward pass.

from trl import SFTConfig

training_args = SFTConfig(..., gradient_checkpointing=True)

Gradient checkpointing diaktifkan secara default di semua trainer untuk mengoptimalkan penggunaan memori. Anda dapat menonaktifkannya dengan mengatur gradient_checkpointing=False jika diperlukan.

Checkpointing aktivasi selektif

Dengan SFTTrainer, Anda dapat menyimpan output perhatian selama forward pass alih-alih menghitungnya ulang pada backward pass. Hal ini memulihkan sebagian besar perlambatan checkpointing pada konteks panjang, dengan satu tensor tambahan berukuran hidden-state per layer. Fitur ini memaksa penggunaan checkpointing non-reentrant.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

Untuk teknik optimasi memori lainnya, lihat Panduan Performa Transformers.