Optimasi 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 memerlukan 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 Kebutuhan 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 memerlukan sekitar 4X GB VRAM dalam presisi 32-bit
• Memuat model dengan X miliar parameter memerlukan sekitar 2X GB VRAM dalam presisi 16-bit bfloat16/float16

Misalnya, memuat model GPT-3 dengan 175 miliar parameter akan memerlukan 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 paralelisme 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.

Mengkuantifikasi Memori Inferensi

Mari kita kuatifikasi kebutuhan 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

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Output:

Penggunaan memori GPU puncak 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 presisi data yang lebih rendah seperti 8-bit integer (int8) atau 4-bit integer 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>

Output:

# 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())

Output:

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 dengan 15 miliar parameter pada GPU konsumen seperti RTX 3090 (24GB VRAM).

Namun, perlu diingat bahwa kuantisasi yang lebih agresif seperti 4-bit dapat kadang-kadang menyebabkan degradasi 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.

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

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

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

Misalnya, fine-tuning model LLaMA dengan 7 miliar parameter akan memerlukan 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 model seluruhnya di seluruh 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 pengoptimasi di seluruh GPU. Ini mengurangi memori dibandingkan dengan paralelisme data klasik dengan hanya menyimpan data partisi yang diperlukan pada setiap GPU selama berbagai fase pelatihan.
3. Paralelisme Tensor: Sebagai gantinya mereplikasi model, paralelisme tensor membagi parameter model menjadi baris atau kolom dan mendistribusikannya di seluruh GPU. Setiap GPU beroperasi pada himpunan parameter, gradien, dan keadaan pengoptimasi yang dipartisi, menghasilkan penghematan memori yang signifikan.
4. Paralelisme Pipa: Teknik ini mempartisi lapisan model di seluruh 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 pengoptimasi bervariasi di seluruh teknik. Selain itu, komponen yang berbeda seperti tubuh transformer dan kepala pemodelan bahasa dapat menampilkan perilaku alokasi memori yang berbeda.

Solusi LLMem

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

Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLM

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

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