Mengoptimalkan Memori untuk Inferensi dan Fine-Tuning Model Bahasa Besar

Model bahasa besar (LLM) seperti GPT-4, Bloom, dan LLaMA telah mencapai kemampuan luar biasa dengan menskalakan hingga miliaran parameter. Namun, menggelar model besar ini untuk inferensi atau fine-tuning sangat menantang karena persyaratan memori yang sangat besar. Dalam blog teknis ini, kita akan menjelajahi teknik untuk memperkirakan dan mengoptimalkan konsumsi memori selama inferensi dan fine-tuning LLM di berbagai pengaturan perangkat keras.

Memahami Persyaratan Memori

Memori yang dibutuhkan untuk memuat LLM sebagian besar ditentukan oleh jumlah parameter dan presisi numerik yang digunakan untuk menyimpan parameter. Aturan sederhana adalah:

• Memuat model dengan X miliar parameter membutuhkan sekitar 4X GB VRAM dalam presisi float 32-bit
• Memuat model dengan X miliar parameter membutuhkan sekitar 2X GB VRAM dalam presisi bfloat16/float16

Misalnya, memuat model GPT-3 dengan 175 miliar parameter akan membutuhkan sekitar 350GB VRAM dalam presisi bfloat16. Saat ini, GPU komersial terbesar seperti NVIDIA A100 dan H100 hanya menawarkan 80GB VRAM, sehingga memerlukan teknik paralelisme tensor dan model.

Selama inferensi, jejak memori didominasi oleh parameter model dan tensor aktivasi sementara yang dihasilkan. Perkiraan tingkat memori puncak selama inferensi adalah jumlah memori yang dibutuhkan untuk memuat parameter model dan memori untuk aktivasi.

Mengukur Memori Inferensi

Mari kita ukur persyaratan memori untuk inferensi menggunakan model OctoCode, yang memiliki sekitar 15 miliar parameter dalam format bfloat16 (~ 31GB). Kita akan menggunakan Transformers library untuk memuat model dan menghasilkan teks:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Pertanyaan: Tolong tulis fungsi Python untuk mengubah byte menjadi gigabyte.\n\nJawaban:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Keluaran:

Penggunaan memori puncak GPU sekitar 29GB, yang sesuai dengan perkiraan kita sebesar 31GB untuk memuat parameter model dalam format bfloat16.

Mengoptimalkan Memori Inferensi dengan Kuantisasi

Sementara bfloat16 adalah presisi yang umum digunakan untuk pelatihan LLM, peneliti telah menemukan bahwa mengkuantifikasi bobot model ke tipe data presisi yang lebih rendah seperti integer 8-bit (int8) atau integer 4-bit dapat secara signifikan mengurangi penggunaan memori dengan kerugian akurasi minimal untuk tugas inferensi seperti generasi teks.

Mari kita lihat penghematan memori dari kuantisasi 8-bit dan 4-bit dari model OctoCode:

</div>
# Kuantisasi 8-bit
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Keluaran:

# Kuantisasi 4-bit
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Keluaran:

Dengan kuantisasi 8-bit, kebutuhan memori turun dari 31GB menjadi 15GB, sedangkan kuantisasi 4-bit lebih lanjut menguranginya menjadi hanya 9,5GB! Ini memungkinkan menjalankan model OctoCode 15 miliar parameter pada GPU konsumen seperti RTX 3090 (24GB VRAM).

Namun, perlu diingat bahwa kuantisasi yang lebih agresif seperti 4-bit terkadang dapat menyebabkan penurunan akurasi dibandingkan dengan presisi 8-bit atau bfloat16. Ada trade-off antara penghematan memori dan akurasi yang harus dievaluasi pengguna untuk kasus penggunaan mereka.

Kuantisasi adalah teknik yang kuat yang dapat memungkinkan penerapan LLM pada lingkungan yang terbatas sumber daya seperti instance cloud, perangkat edge, atau bahkan ponsel dengan mengurangi jejak memori secara drastis.

Mengestimasi Memori untuk Fine-Tuning

Sementara kuantisasi sebagian besar digunakan untuk inferensi yang efisien, teknik seperti paralelisme tensor dan paralelisme model sangat penting untuk mengelola kebutuhan memori selama pelatihan atau fine-tuning model bahasa besar.

Konsumsi memori puncak selama fine-tuning biasanya 3-4 kali lebih tinggi daripada inferensi karena kebutuhan memori tambahan untuk:

• Gradien
• Keadaan pengoptimal
• Aktivasi dari jalur maju yang disimpan untuk backpropagation

Perkiraan konservatif adalah bahwa fine-tuning LLM dengan X miliar parameter membutuhkan sekitar 4 * (2X) = 8X GB VRAM dalam presisi bfloat16.

Misalnya, fine-tuning model LLaMA 7 miliar parameter akan membutuhkan sekitar 7 * 8 = 56GB VRAM per GPU dalam presisi bfloat16. Ini melebihi kapasitas memori GPU saat ini, sehingga memerlukan teknik fine-tuning terdistribusi.

Teknik Fine-Tuning Terdistribusi

Beberapa metode fine-tuning terdistribusi telah diajukan untuk mengatasi keterbatasan memori GPU untuk model besar:

1. Paralelisme Data: Pendekatan paralelisme data klasik mereplikasi seluruh model di beberapa GPU sambil membagi dan mendistribusikan batch data pelatihan. Ini mengurangi waktu pelatihan secara linear dengan jumlah GPU tetapi tidak mengurangi kebutuhan memori puncak pada setiap GPU.
2. ZeRO Stage 3: Bentuk paralelisme data yang lebih maju yang mempartisi parameter model, gradien, dan keadaan pengoptimal di beberapa GPU. Ini mengurangi memori dibandingkan dengan paralelisme data klasik dengan menyimpan hanya data partisi yang diperlukan pada setiap GPU selama fase pelatihan yang berbeda.
3. Paralelisme Tensor: Sebagai gantinya mereplikasi model, paralelisme tensor membagi parameter model menjadi baris atau kolom dan mendistribusikannya di beberapa GPU. Setiap GPU beroperasi pada himpunan parameter, gradien, dan keadaan pengoptimal yang dipartisi, menghasilkan penghematan memori yang signifikan.
4. Paralelisme Pipa: Teknik ini mempartisi lapisan model di beberapa GPU/pekerja, dengan setiap perangkat menjalankan subset lapisan. Aktivasi dilewatkan antara pekerja, mengurangi memori puncak tetapi meningkatkan overhead komunikasi.

Mengestimasi penggunaan memori untuk metode terdistribusi ini tidaklah sederhana karena distribusi parameter, gradien, aktivasi, dan keadaan pengoptimal bervariasi di antara teknik. Selain itu, komponen yang berbeda seperti tubuh transformer dan kepala pemodelan bahasa mungkin menunjukkan perilaku alokasi memori yang berbeda.

Solusi LLMem

Peneliti baru-baru ini mengusulkan LLMem, sebuah solusi yang secara akurat memperkirakan konsumsi memori GPU saat menerapkan metode fine-tuning terdistribusi ke LLM di beberapa GPU.

Mengestimasi Penggunaan Memori GPU untuk Fine-Tuning LLM Pra-Pelatihan

LLMem mempertimbangkan faktor-faktor seperti menggabungkan kembali parameter sebelum komputasi (ZeRO Stage 3), pengumpulan output dalam jalur mundur (paralelisme tensor), dan strategi alokasi memori yang berbeda untuk tubuh transformer dan kepala pemodelan bahasa.

Hasil eksperimen menunjukkan bahwa LLMem dapat memperkirakan penggunaan memori GPU puncak untuk fine-tuning LLM pada satu GPU dengan tingkat kesalahan hingga 1,6%, mengungguli rata-rata kesalahan DNNMem sebesar 42,6%. Ketika menerapkan metode fine-tuning terdistribusi ke LLM dengan lebih dari satu miliar parameter pada beberapa GPU, LLMem mencapai rata-rata kesalahan 3,0%.