Mengoptimalkan LLM untuk Kecepatan dan Memori

Model Bahasa Besar (LLM) seperti GPT3/4, Falcon, dan Llama berkembang pesat dalam kemampuannya untuk menangani tugas-tugas yang berpusat pada manusia, sehingga menjadi alat penting di industri berbasis pengetahuan modern. Namun, penerapan model-model ini dalam tugas dunia nyata masih menantang:

Untuk menunjukkan kemampuan pemahaman dan generasi teks yang mendekati manusia, LLM saat ini harus terdiri dari miliaran parameter (lihat Kaplan et al, Wei et. al). Hal ini pada akhirnya meningkatkan kebutuhan memori untuk inferensi.
Dalam banyak tugas dunia nyata, LLM perlu diberikan informasi kontekstual yang luas. Hal ini menuntut kemampuan model untuk mengelola urutan input yang sangat panjang selama inferensi.

Inti dari tantangan-tantangan ini terletak pada peningkatan kemampuan komputasi dan memori LLM, terutama saat menangani sekuens input yang luas.

Dalam panduan ini, kita akan membahas teknik-teknik efektif untuk deployment LLM yang efisien:

Presisi Lebih Rendah: Penelitian menunjukkan bahwa beroperasi pada presisi numerik yang lebih rendah, yaitu 8-bit dan 4-bit, dapat memberikan keunggulan komputasi tanpa penurunan kinerja model yang signifikan.
Flash Attention: Flash Attention adalah variasi dari algoritma attention yang tidak hanya menyediakan pendekatan yang lebih hemat memori, tetapi juga mewujudkan peningkatan efisiensi berkat pemanfaatan memori GPU yang dioptimalkan.
Inovasi Arsitektur: Mengingat LLM selalu di-deploy dengan cara yang sama selama inferensi, yaitu generasi teks autoregresif dengan konteks input yang panjang, arsitektur model khusus telah diusulkan yang memungkinkan inferensi yang lebih efisien. Kemajuan terpenting dalam arsitektur model di sini adalah Alibi, Rotary embeddings, Multi-Query Attention (MQA), dan Grouped-Query-Attention (GQA).

Sepanjang panduan ini, kami akan menyajikan analisis tentang generasi auto-regresif dari perspektif tensor. Kami membahas kelebihan dan kekurangan adopsi presisi yang lebih rendah, memberikan eksplorasi komprehensif tentang algoritma attention terbaru, serta mendiskusikan arsitektur LLM yang lebih baik. Selama proses tersebut, kami menjalankan contoh praktis yang menampilkan masing-masing peningkatan fitur.

1. Presisi Lebih Rendah

Kebutuhan memori LLM dapat dipahami dengan terbaiknya dengan memandang LLM sebagai seperangkat matriks dan vektor bobot, serta input teks sebagai rangkaian vektor. Selanjutnya, definisi bobot akan digunakan untuk merujuk pada semua matriks dan vektor bobot model.

Pada saat panduan ini ditulis, LLM terdiri dari setidaknya beberapa miliar parameter. Setiap parameter terdiri dari sebuah bilangan desimal, misalnya 4.5689 yang biasanya disimpan dalam format float32, bfloat16, atau float16. Hal ini memungkinkan kita untuk dengan mudah menghitung kebutuhan memori untuk memuat LLM ke dalam memori:

Memuat bobot model dengan X miliar parameter membutuhkan sekitar 4 * X GB VRAM dalam presisi float32

Namun, saat ini model jarang dilatih dalam presisi penuh float32, melainkan biasanya dalam presisi bfloat16 atau, lebih jarang, dalam presisi float16. Oleh karena itu, aturan praktisnya menjadi:

Memuat bobot model dengan X miliar parameter membutuhkan sekitar 2 * X GB VRAM dalam presisi bfloat16/float16

Untuk input teks yang lebih pendek (kurang dari 1024 token), kebutuhan memori untuk inferensi sangat didominasi oleh kebutuhan memori untuk memuat bobot. Oleh karena itu, untuk saat ini, mari kita asumsikan bahwa kebutuhan memori untuk inferensi sama dengan kebutuhan memori untuk memuat model ke VRAM GPU.

Sebagai contoh, berikut perkiraan berapa banyak VRAM yang dibutuhkan untuk memuat model dalam bfloat16:

GPT3 membutuhkan 2 * 175 GB = 350 GB VRAM
Bloom membutuhkan 2 * 176 GB = 352 GB VRAM
Llama-2-70b membutuhkan 2 * 70 GB = 140 GB VRAM
Falcon-40b membutuhkan 2 * 40 GB = 80 GB VRAM
MPT-30b membutuhkan 2 * 30 GB = 60 GB VRAM
bigcode/starcoder membutuhkan 2 * 15,5 = 31 GB VRAM

Hingga dokumen ini ditulis, chip GPU terbesar di pasar adalah A100 & H100 yang menawarkan 80GB VRAM. Sebagian besar model yang disebutkan sebelumnya membutuhkan lebih dari 80GB hanya untuk dimuat, sehingga secara pasti memerlukan parallelisme tensor dan/atau parallelisme pipeline.

🤗 Transformers kini mendukung tensor parallelism untuk model-model yang didukung yang memiliki base_tp_plan di kelas config masing-masing. Pelajari lebih lanjut tentang Tensor Parallelism di sini. Selain itu, jika Anda tertarik menulis model dengan cara yang ramah tensor parallelism, silakan lihat library text-generation-inference.

Pipeline parallelism naif didukung secara bawaan. Untuk ini, cukup muat model dengan device="auto" yang akan secara otomatis menempatkan lapisan yang berbeda pada GPU yang tersedia seperti dijelaskan di sini. Perlu dicatat, meskipun sangat efektif, pipeline parallelism naif ini tidak mengatasi masalah GPU yang menganggur. Untuk ini, pipeline parallelism yang lebih canggih diperlukan seperti dijelaskan di sini.

Jika Anda memiliki akses ke node 8 x 80GB A100, Anda dapat memuat BLOOM sebagai berikut

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

Dengan menggunakan device_map="auto", lapisan perhatian akan didistribusikan secara merata di semua GPU yang tersedia.

Dalam panduan ini, kami akan menggunakan bigcode/octocoder karena model ini dapat dijalankan pada satu perangkat GPU A100 dengan memori 40 GB. Perhatikan bahwa semua optimasi memori dan kecepatan yang akan kami terapkan selanjutnya juga berlaku sama untuk model yang memerlukan model atau tensor parallelism.

Karena model dimuat dalam presisi bfloat16, menggunakan aturan praktis di atas, kita dapat mengharapkan kebutuhan memori untuk menjalankan inferensi dengan bigcode/octocoder sekitar 31 GB VRAM. Mari kita coba.

Pertama, kita memuat model dan tokenizer, lalu meneruskan keduanya ke objek pipeline dari Transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", 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 function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Output:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Bagus, kita sekarang dapat langsung menggunakan hasilnya untuk mengonversi byte menjadi Gigabyte.

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

Mari kita panggil torch.accelerator.memory.max_memory_allocated untuk mengukur alokasi memori puncak akselerator.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Output:

29.0260648727417

Cukup dekat dengan perhitungan kasar kami! Kita bisa melihat bahwa angkanya tidak persis benar karena konversi dari byte ke kilobyte memerlukan perkalian dengan 1024, bukan 1000. Oleh karena itu, rumus kasar ini juga dapat dipahami sebagai perhitungan “maksimal X GB”. Perlu dicatat bahwa jika kita mencoba menjalankan model dalam presisi float32 penuh, dibutuhkan VRAM sebesar 64 GB.

Hampir semua model saat ini dilatih dalam bfloat16, sehingga tidak ada alasan untuk menjalankan model dalam presisi penuh float32 jika GPU Anda mendukung bfloat16. Float32 tidak akan memberikan hasil inferensi yang lebih baik daripada presisi yang digunakan untuk melatih model.

Jika Anda tidak yakin dalam format apa bobot model disimpan di Hub, Anda selalu dapat memeriksa konfigurasi checkpoint di bawah "dtype", misalnya di sini. Disarankan untuk mengatur model ke tipe presisi yang sama seperti yang tertulis di konfigurasi saat memuat dengan from_pretrained(..., dtype=...), kecuali jika tipe aslinya adalah float32, di mana Anda dapat menggunakan baik float16 atau bfloat16 untuk inferensi.

Mari kita definisikan fungsi flush(...) untuk membebaskan semua memori yang dialokasikan sehingga kita dapat mengukur memori GPU puncak yang dialokasikan secara akurat.

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

Mari kita putuskan sekarang untuk eksperimen berikutnya.

flush()

Dari pustaka Accelerate, Anda juga dapat menggunakan metode utilitas yang tidak bergantung pada perangkat bernama release_memory, yang mempertimbangkan berbagai backend perangkat keras seperti XPU, MLU, NPU, MPS, dan lainnya.

from accelerate.utils import release_memory
# ...

release_memory(model)

Bagaimana jika GPU Anda tidak memiliki 32 GB VRAM? Telah ditemukan bahwa bobot model dapat dikuantisasi menjadi 8-bit atau 4-bit tanpa penurunan kinerja yang signifikan (lihat Dettmers et al.). Model bahkan dapat dikuantisasi menjadi 3 atau 2 bit dengan penurunan kinerja yang masih dapat diterima, seperti yang ditunjukkan dalam makalah GPTQ terbaru 🤯.

Tanpa masuk ke terlalu banyak detail, skema kuantisasi bertujuan untuk mengurangi presisi bobot sambil tetap menjaga hasil inferensi model agar seakurat mungkin (alias sedekat mungkin dengan bfloat16). Perlu dicatat bahwa kuantisasi bekerja sangat baik untuk generasi teks karena yang kita pedulikan hanyalah memilih set token berikutnya yang paling mungkin dan tidak terlalu peduli dengan nilai eksak dari distribusi logit token berikutnya. Yang penting adalah distribusi logit token berikutnya tetap kurang lebih sama sehingga operasi argmax atau topk memberikan hasil yang sama.

Terdapat berbagai teknik kuantisasi, yang tidak akan kami bahas secara detail di sini, tetapi secara umum, semua teknik kuantisasi bekerja sebagai berikut:

1. Kuantisasi semua bobot ke presisi target
1. Muat bobot terkuantisasi, dan lewati urutan input vektor dalam presisi bfloat16
1. Dequantisasi bobot secara dinamis ke bfloat16 untuk melakukan komputasi dengan vektor input mereka dalam presisi bfloat16

Singkatnya, ini berarti bahwa matriks perkalian input dengan $X$ adalah input, $W$ adalah matriks bobot dan $Y$ adalah output: $Y = X * W$

diubah menjadi $Y = X * \text{dequantize}(W)$

untuk setiap perkalian matriks. Dequantization dan re-quantization dilakukan secara berurutan untuk semua matriks bobot saat input berjalan melalui graf jaringan.

Oleh karena itu, waktu inferensi sering kali tidak berkurang saat menggunakan bobot terkuantisasi, melainkan justru meningkat. Cukup teori, mari kita coba! Untuk mengkuantisasi bobot dengan Transformers, Anda perlu memastikan bahwa library bitsandbytes sudah terinstal.

!pip install bitsandbytes

Kemudian, kita dapat memuat model dalam kuantisasi 8-bit dengan cukup menambahkan flag load_in_8bit=True ke from_pretrained.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

Sekarang, jalankan contoh kita sekali lagi dan ukur penggunaan memori.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Output:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Bagus, kita mendapatkan hasil yang sama seperti sebelumnya, jadi tidak ada penurunan akurasi! Mari kita lihat berapa banyak memori yang digunakan kali ini.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Output:

15.219234466552734

Jauh lebih sedikit! Sekarang hanya sedikit di atas 15 GB, sehingga model ini dapat dijalankan pada GPU konsumen seperti 4090. Kami melihat peningkatan efisiensi memori yang sangat baik dan hampir tidak ada degradasi pada output model. Namun, kami juga dapat mengamati perlambatan ringan selama inferensi.

Kami menghapus model dan mengosongkan memori sekali lagi.

del model
del pipe

flush()

Mari kita lihat berapa konsumsi memori GPU puncak yang dihasilkan oleh kuantisasi 4-bit. Mengkuantisasi model ke 4-bit dapat dilakukan dengan API yang sama seperti sebelumnya - kali ini dengan meneruskan load_in_4bit=True alih-alih load_in_8bit=True.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=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):]
result

Output:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

Kita hampir melihat teks output yang sama seperti sebelumnya - hanya python yang hilang tepat sebelum cuplikan kode. Mari kita lihat berapa banyak memori yang dibutuhkan.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Output:

9.543574333190918

Hanya 9,5GB! Itu benar-benar tidak banyak untuk model dengan parameter >15 miliar.

Meskipun kami melihat penurunan akurasi yang sangat kecil pada model kami di sini, kuantisasi 4-bit dalam praktiknya sering kali dapat menghasilkan hasil yang berbeda dibandingkan dengan kuantisasi 8-bit atau inferensi bfloat16 penuh. Terserah pengguna untuk mencobanya.

Perhatikan juga bahwa inferensi di sini sedikit lebih lambat dibandingkan kuantisasi 8-bit, yang disebabkan oleh metode kuantisasi yang lebih agresif yang digunakan untuk kuantisasi 4-bit, sehingga $\text{quantize}$ dan $\text{dequantize}$ membutuhkan waktu lebih lama selama inferensi.

del model
del pipe

flush()

Secara keseluruhan, kami melihat bahwa menjalankan OctoCoder dalam presisi 8-bit mengurangi VRAM GPU yang dibutuhkan dari 32G VRAM GPU menjadi hanya 15GB, dan menjalankan model dalam presisi 4-bit lebih lanjut mengurangi VRAM GPU yang dibutuhkan menjadi hanya sedikit di atas 9GB.

Kuantisasi 4-bit memungkinkan model dijalankan pada GPU seperti RTX3090, V100, dan T4 yang cukup mudah diakses oleh kebanyakan orang.

Untuk informasi lebih lanjut tentang kuantisasi dan untuk melihat bagaimana seseorang dapat mengkuantisasi model agar membutuhkan memori GPU VRAM bahkan lebih sedikit dari 4-bit, kami merekomendasikan untuk melihat implementasi GPT-QModel.

Sebagai kesimpulan, penting untuk diingat bahwa kuantisasi model menukar efisiensi memori yang lebih baik dengan akurasi dan, dalam beberapa kasus, waktu inferensi.

Jika memori GPU bukan kendala untuk kasus penggunaan Anda, seringkali tidak perlu mencari tahu tentang kuantisasi. Namun, banyak GPU yang tidak dapat menjalankan LLM tanpa metode kuantisasi, dan dalam kasus ini, skema kuantisasi 4-bit dan 8-bit adalah alat yang sangat berguna.

Untuk informasi penggunaan yang lebih rinci, kami sangat menyarankan untuk melihat Dokumentasi Kuantisasi Transformers. Selanjutnya, mari kita lihat bagaimana kita dapat meningkatkan efisiensi komputasi dan memori dengan menggunakan algoritma yang lebih baik dan arsitektur model yang ditingkatkan.

2. Flash Attention

LLM berkinerja terbaik saat ini pada dasarnya berbagi arsitektur fundamental yang kurang lebih sama, yang terdiri dari lapisan feed-forward, lapisan aktivasi, lapisan normalisasi lapisan, dan yang paling krusial, lapisan self-attention.

Lapisan self-attention merupakan inti dari Large Language Models (LLM) karena memungkinkan model memahami hubungan kontekstual antar token input. Namun, konsumsi memori GPU puncak untuk lapisan self-attention meningkat secara kuadratik baik dalam kompleksitas komputasi maupun memori seiring dengan jumlah token input (juga disebut panjang sekuens) yang kami notasikan berikut dengan $N$ . Meskipun hal ini tidak terlalu terasa untuk sekuens input yang lebih pendek (hingga 1000 token input), masalah ini menjadi serius untuk sekuens input yang lebih panjang (sekitar 16000 token input).

Mari kita lihat lebih dekat. Rumus untuk menghitung output $\mathbf{O}$ dari lapisan self-attention untuk input $\mathbf{X}$ dengan panjang $N$ adalah: $\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$ adalah urutan input untuk lapisan attention. Proyeksi $\mathbf{Q}$ dan $\mathbf{K}$ masing-masing akan terdiri dari $N$ vektor sehingga $\mathbf{QK}^T$ berukuran $N^2$ .

LLM biasanya memiliki beberapa attention head, sehingga melakukan beberapa komputasi self-attention secara paralel. Dengan asumsi LLM memiliki 40 attention head dan berjalan dalam presisi bfloat16, kita dapat menghitung kebutuhan memori untuk menyimpan $\mathbf{QK^T}$ matriks menjadi $40 * 2 * N^2$ byte. Untuk $N=1000$ hanya sekitar 50 MB VRAM yang dibutuhkan, namun untuk $N=16000$ kita akan membutuhkan 19 GB VRAM, dan untuk $N=100,000$ kita akan membutuhkan hampir 1 TB hanya untuk menyimpan $\mathbf{QK}^T$ matriks.

Singkatnya, algoritma self-attention default dengan cepat menjadi sangat mahal dari segi memori untuk konteks input yang besar.

Seiring dengan peningkatan kemampuan LLM dalam pemahaman dan generasi teks, model-model ini diterapkan pada tugas-tugas yang semakin kompleks. Jika sebelumnya model hanya menangani terjemahan atau ringkasan beberapa kalimat, kini mereka mengelola seluruh halaman, yang menuntut kemampuan untuk memproses panjang input yang sangat besar.

Bagaimana kita bisa menghilangkan kebutuhan memori yang sangat besar untuk panjang input yang besar? Kita membutuhkan cara baru untuk menghitung mekanisme self-attention yang menghilangkan $\mathbf{QK}^T$ matriks. Tri Dao dkk. mengembangkan algoritma baru yang tepat seperti itu dan menamainya Flash Attention.

Singkatnya, Flash Attention memisahkan $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ perhitungan dan sebaliknya menghitung potongan-potongan kecil output dengan mengulang beberapa langkah perhitungan softmax: $\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$

dengan $s^a_{ij}$ dan $s^b_{ij}$ merupakan beberapa statistik normalisasi softmax yang perlu dihitung ulang untuk setiap $i$ dan $j$ .

Perlu dicatat bahwa keseluruhan Flash Attention sedikit lebih kompleks dan sangat disederhanakan di sini karena penjelasan yang terlalu mendalam berada di luar cakupan panduan ini. Pembaca diundang untuk melihat makalah Flash Attention yang ditulis dengan baik untuk detail lebih lanjut.

Poin utama di sini adalah:

Dengan melacak statistik normalisasi softmax dan menggunakan beberapa matematika cerdas, Flash Attention memberikan output yang identik secara numerik dibandingkan dengan lapisan self-attention default, dengan biaya memori yang hanya meningkat secara linear terhadap $N$ .

Jika melihat rumusnya, secara intuitif orang akan mengatakan bahwa Flash Attention pasti jauh lebih lambat dibandingkan rumus self-attention default karena lebih banyak komputasi yang harus dilakukan. Memang, Flash Attention membutuhkan lebih banyak FLOPs dibandingkan attention normal karena statistik normalisasi softmax harus terus-menerus dihitung ulang (lihat paper untuk detail lebih lanjut jika tertarik)

Namun, Flash Attention jauh lebih cepat dalam inferensi dibandingkan perhatian default, berkat kemampuannya untuk secara signifikan mengurangi tuntutan pada memori GPU yang lebih lambat dan berbandwidth tinggi (VRAM), dengan berfokus pada memori on-chip yang lebih cepat (SRAM).

Pada dasarnya, Flash Attention memastikan bahwa semua operasi tulis dan baca perantara dapat dilakukan menggunakan memori on-chip SRAM yang cepat, alih-alih harus mengakses memori VRAM yang lebih lambat untuk menghitung vektor output $\mathbf{O}$ .

Dalam praktiknya, saat ini sama sekali tidak ada alasan untuk tidak menggunakan Flash Attention jika tersedia. Algoritma ini menghasilkan output yang secara matematis sama, serta lebih cepat dan lebih hemat memori.

3. Inovasi Arsitektural

Sampai sejauh ini, kita telah membahas peningkatan efisiensi komputasi dan memori dengan cara:

Mencasting bobot ke format presisi yang lebih rendah
Mengganti algoritma self-attention dengan versi yang lebih hemat memori dan komputasi

Sekarang, mari kita lihat bagaimana kita dapat mengubah arsitektur LLM agar paling efektif dan efisien untuk tugas yang memerlukan input teks panjang, misalnya:

Pertanyaan dan Jawaban yang Diperkaya dengan Pemulihan,
Ringkasan,
Chat

Perhatikan bahwa chat tidak hanya mengharuskan LLM untuk menangani input teks panjang, tetapi juga menuntut agar LLM mampu menangani dialog bolak-balik antara pengguna dan asisten (seperti ChatGPT) secara efisien.

Setelah dilatih, arsitektur LLM mendasar sulit untuk diubah, sehingga penting untuk mempertimbangkan tugas-tugas LLM sebelumnya dan mengoptimalkan arsitektur model secara sesuai. Ada dua komponen penting dalam arsitektur model yang dengan cepat menjadi hambatan memori dan/atau kinerja untuk urutan input yang besar.

Embedding posisi
Cache kunci-nilai

Mari kita bahas setiap komponen secara lebih detail

3.1 Meningkatkan embedding posisi LLM

Self-attention menempatkan setiap token dalam hubungan dengan token lainnya. Sebagai contoh, matriks $\text{Softmax}(\mathbf{QK}^T)$ dari urutan teks input ““Hello”, “I”, “love”, “you” bisa terlihat seperti berikut:

Setiap token kata diberikan massa probabilitas yang menunjukkan seberapa besar perhatiannya terhadap semua token kata lainnya, sehingga token tersebut ditempatkan dalam relasi dengan semua token kata lainnya. Misalnya, kata “love” memperhatikan kata “Hello” sebesar 5%, “I” sebesar 30%, dan dirinya sendiri sebesar 65%.

LLM yang berbasis self-attention, tetapi tanpa position embeddings akan mengalami kesulitan besar dalam memahami posisi input teks satu sama lain. Hal ini karena skor probabilitas yang dihitung oleh $\mathbf{QK}^T$ mengaitkan setiap token kata dengan setiap token kata lainnya dalam $O(1)$ perhitungan, terlepas dari jarak posisi relatif mereka satu sama lain. Oleh karena itu, bagi LLM tanpa position embeddings, setiap token tampaknya memiliki jarak yang sama ke semua token lainnya, e.g. membedakan antara “Hello I love you” dan “You love I hello” akan menjadi sangat menantang.

Agar LLM dapat memahami urutan kalimat, diperlukan petunjuk tambahan yang biasanya diterapkan dalam bentuk positional encodings (atau juga disebut positional embeddings). Positional encodings mengodekan posisi setiap token ke dalam representasi numerik yang dapat dimanfaatkan oleh LLM untuk lebih memahami urutan kalimat.

Para penulis Attention Is All You Need mengkenalkan positional embedding sinusoidal $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$ . di mana setiap vektor $\mathbf{p}_i$ dihitung sebagai fungsi sinusoidal dari posisinya $i$ . Posisi encoding kemudian ditambahkan secara sederhana ke vektor urutan input $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$ sehingga memberi petunjuk kepada model untuk lebih baik mempelajari urutan kalimat.

Alih-alih menggunakan position embeddings tetap, peneliti lain (seperti Devlin et al.) menggunakan positional encodings yang dipelajari, di mana positional embeddings $\mathbf{P}$ dipelajari selama pelatihan.

Embedding posisi sinusoidal dan yang dipelajari (learned) dulunya menjadi metode utama untuk mengkodekan urutan kalimat ke dalam LLM, namun beberapa masalah terkait encoding posisi ini ditemukan:

Sinusoidal dan learned position embeddings keduanya merupakan absolute positional embeddings, i.e. yaitu mengkodekan embedding unik untuk setiap position id: $0, \ldots, N$ . Seperti yang ditunjukkan oleh Huang et al. dan Su et al., absolute positional embeddings menyebabkan kinerja LLM yang buruk untuk input teks panjang. Untuk input teks panjang, lebih menguntungkan jika model mempelajari jarak posisi relatif antar token input daripada posisi absolut mereka.
Saat menggunakan learned position embeddings, LLM harus dilatih pada panjang input tetap $N$, yang membuatnya sulit untuk mengekstrapolasi ke panjang input yang lebih panjang dari yang digunakan saat pelatihan.

Baru-baru ini, embedding posisi relatif yang dapat mengatasi masalah yang disebutkan di atas menjadi semakin populer, terutama:

Keduanya RoPE dan ALiBi berargumen bahwa cara terbaik untuk memberi isyarat urutan kalimat kepada LLM adalah secara langsung di dalam algoritma self-attention, karena di situlah token kata saling berhubungan. Secara lebih spesifik, urutan kalimat harus diberi isyarat dengan memodifikasi $\mathbf{QK}^T$ perhitungan.

Tanpa masuk ke terlalu banyak detail, RoPE menyatakan bahwa informasi posisi dapat dikodekan ke dalam pasangan query-key, e.g. $\mathbf{q}_i$ dan $\mathbf{x}_j$ dengan memutar setiap vektor sebesar sudut $\theta * i$ dan $\theta * j$ masing-masing dengan $i, j$ yang menggambarkan posisi kalimat dari masing-masing vektor: $\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$

$\mathbf{R}_{\theta, i - j}$ dengan demikian merepresentasikan matriks rotasi. $\theta$ adalah tidak dipelajari selama pelatihan, melainkan diatur ke nilai yang telah ditentukan sebelumnya yang bergantung pada panjang urutan input maksimum selama pelatihan.

Dengan cara ini, skor probabilitas antara $\mathbf{q}_i$ dan $\mathbf{q}_j$ hanya terpengaruh jika $i \ne j$ dan sepenuhnya bergantung pada jarak relatif $i - j$ tanpa memperhatikan posisi spesifik masing-masing vektor $i$ dan $j$ .

RoPE digunakan dalam beberapa LLM terpenting saat ini, seperti:

Sebagai alternatif, ALiBi mengusulkan skema pengkodean posisi relatif yang jauh lebih sederhana. Jarak relatif antar token input ditambahkan sebagai bilangan bulat negatif yang diskalakan dengan nilai yang telah ditentukan sebelumnya m ke setiap entri query-key dari $\mathbf{QK}^T$ matriks tepat sebelum komputasi softmax.

Seperti yang ditunjukkan dalam paper ALiBi, encoding posisi relatif yang sederhana ini memungkinkan model untuk mempertahankan kinerja tinggi bahkan pada urutan input teks yang sangat panjang.

ALiBi digunakan dalam beberapa LLM terpenting saat ini, seperti:

Keduanya RoPE dan ALiBi position encodings dapat diekstrapolasi ke panjang input yang tidak terlihat selama pelatihan, sementara telah ditunjukkan bahwa ekstrapolasi bekerja jauh lebih baik secara bawaan untuk ALiBi dibandingkan dengan RoPE. Untuk ALiBi, cukup meningkatkan nilai matriks posisi segitiga bawah agar sesuai dengan panjang urutan input. Untuk RoPE, menjaga $\theta$ yang digunakan selama pelatihan menghasilkan hasil yang buruk saat memasukkan teks yang jauh lebih panjang dari yang dilihat selama pelatihan, c.f Press et al.. Namun, komunitas telah menemukan beberapa trik efektif yang menyesuaikan $\theta$, sehingga memungkinkan RoPE position embeddings bekerja dengan baik untuk urutan input teks yang diekstrapolasi (lihat di sini).

Baik RoPE maupun ALiBi adalah embedding posisi relatif yang tidak dipelajari selama pelatihan, melainkan didasarkan pada intuisi berikut:

Petunjuk posisi tentang input teks harus diberikan langsung ke $\mathbf{QK}^T$ matriks dari lapisan self-attention.
LLM harus diberi insentif untuk mempelajari encoding posisi jarak relatif yang konstan.
Semakin jauh jarak antar token input teks, semakin rendah probabilitas query-value mereka. Baik RoPE maupun ALiBi menurunkan probabilitas query-key untuk token yang berada jauh satu sama lain. RoPE melakukannya dengan mengurangi hasil perkalian vektor mereka dengan memperbesar sudut antara vektor query-key. ALiBi melakukannya dengan menambahkan angka negatif besar ke hasil perkalian vektor.

Kesimpulannya, LLM yang ditujukan untuk tugas yang memerlukan penanganan input teks besar lebih baik dilatih dengan positional embedding relatif, seperti RoPE dan ALiBi. Perhatikan juga bahwa meskipun LLM dengan RoPE dan ALiBi hanya dilatih pada panjang tetap, misalnya $N_1 = 2048$ ia masih dapat digunakan dalam praktik dengan input teks yang jauh lebih besar dari $N_1$, seperti $N_2 = 8192 > N_1$, dengan melakukan ekstrapolasi pada positional embedding.

3.2 Cache kunci-nilai

Generasi teks auto-regresif dengan LLM bekerja dengan cara memasukkan urutan input secara iteratif, mengambil sampel token berikutnya, menambahkan token berikutnya ke urutan input, dan terus melakukannya hingga LLM menghasilkan token yang menandakan bahwa generasi telah selesai.

Silakan lihat Tutorial Generate Text Transformers untuk penjelasan visual yang lebih baik tentang cara kerja generasi auto-regresif.

Mari kita jalankan potongan kode singkat untuk menunjukkan bagaimana auto-regressive bekerja dalam praktiknya. Kita akan mengambil token berikutnya yang paling mungkin melalui torch.argmax.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

Output:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

Seperti yang dapat kita lihat, setiap kali kita meningkatkan token input teks dengan token yang baru saja disampelkan.

Dengan sangat sedikit pengecualian, LLM dilatih menggunakan tujuan pemodelan bahasa kausal dan oleh karena itu memmask matriks segitiga atas dari skor perhatian - inilah mengapa pada dua diagram di atas skor perhatian dibiarkan kosong (alias memiliki probabilitas 0). Untuk ringkasan singkat tentang pemodelan bahasa kausal, Anda dapat merujuk pada blog Illustrated Self Attention.

Sebagai konsekuensinya, token tidak pernah bergantung pada token yang lebih akhir, lebih khususnya vektor $\mathbf{q}_i$ tidak pernah dipasangkan dengan vektor key atau values $\mathbf{k}_j, \mathbf{v}_j$ jika $j > i$. Sebaliknya, $\mathbf{q}_i$ hanya memperhatikan vektor key-value sebelumnya $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , untuk } m \in {0, \ldots i - 1}$. Untuk mengurangi komputasi yang tidak perlu, vektor key-value setiap lapisan untuk semua timestep sebelumnya dapat di-cache.

Berikutnya, kami akan menginstruksikan LLM untuk memanfaatkan cache key-value dengan mengambil dan meneruskannya untuk setiap forward pass. Di Transformers, kami dapat mengambil cache key-value dengan meneruskan flag use_cache ke panggilan forward dan kemudian meneruskannya bersama token saat ini.

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

Output:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

Seperti yang dapat dilihat, ketika menggunakan cache kunci-nilai, token input teks tidak bertambah panjang, melainkan tetap berupa satu vektor input. Di sisi lain, panjang cache kunci-nilai bertambah satu pada setiap langkah decoding.

Pemanfaatan key-value cache berarti bahwa $\mathbf{QK}^T$ pada dasarnya direduksi menjadi $\mathbf{q}_c\mathbf{K}^T$ dengan $\mathbf{q}_c$ sebagai proyeksi query dari token input yang sedang diproses, yang selalu hanya berupa satu vektor.

Menggunakan cache kunci-nilai memiliki dua keuntungan:

Peningkatan efisiensi komputasi yang signifikan karena jumlah komputasi yang dilakukan lebih sedikit dibandingkan dengan menghitung matriks $\mathbf{QK}^T$ secara penuh. Hal ini menyebabkan peningkatan kecepatan inferensi
Memori maksimum yang dibutuhkan tidak meningkat secara kuadratik seiring dengan jumlah token yang dihasilkan, melainkan hanya meningkat secara linear.

Kita selalu harus memanfaatkan key-value cache karena menghasilkan hasil yang identik dan percepatan signifikan untuk urutan input yang lebih panjang. Transformers mengaktifkan key-value cache secara default saat menggunakan text pipeline atau metode generate. Kami memiliki panduan khusus tentang cache di sini.

Perlu dicatat bahwa, terlepas dari saran kami untuk menggunakan key-value caches, output LLM Anda mungkin sedikit berbeda saat menggunakannya. Ini adalah sifat dari kernel perkalian matriks itu sendiri — Anda dapat membaca lebih lanjut tentang hal ini di sini.

3.2.1 Percakapan multi-putaran

Cache key-value sangat berguna untuk aplikasi seperti chat yang memerlukan beberapa langkah decoding auto-regresif. Mari kita lihat sebuah contoh.

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

Dalam obrolan ini, LLM menjalankan dekoding auto-regresif dua kali:

Pada pertama kali, cache kunci-nilai kosong dan prompt input adalah "User: How many people live in France?" dan model menghasilkan teks "Roughly 75 million people live in France" secara auto-regresif sambil meningkatkan cache kunci-nilai di setiap langkah decoding.
Kali kedua, prompt input adalah "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?". Berkat cache, semua vektor key-value untuk dua kalimat pertama sudah dihitung. Oleh karena itu, prompt input hanya terdiri dari "User: And how many in Germany?". Saat memproses prompt input yang dipendekkan, vektor key-value yang dihitung digabungkan ke cache key-value dari decoding pertama. Jawaban Assistant kedua "Germany has ca. 81 million inhabitants" kemudian dihasilkan secara auto-regresif dengan cache key-value yang terdiri dari vektor key-value terenkripsi dari "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?".

Ada dua hal yang perlu diperhatikan di sini:

Menjaga seluruh konteks sangat penting bagi LLM yang diimplementasikan dalam chat agar LLM memahami semua konteks sebelumnya dari percakapan. Misalnya, untuk contoh di atas, LLM perlu memahami bahwa pengguna merujuk pada populasi saat bertanya "And how many are in Germany".
Cache key-value sangat berguna untuk chat karena memungkinkan kita untuk terus memperluas riwayat chat yang telah di-encode, alih-alih harus meng-encode ulang riwayat chat dari awal (seperti yang terjadi, misalnya, saat menggunakan arsitektur encoder-decoder).

Di transformers, pemanggilan generate akan mengembalikan past_key_values ketika return_dict_in_generate=True diteruskan, di samping use_cache=True yang menjadi default. Perhatikan bahwa fitur ini belum tersedia melalui antarmuka pipeline.

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

Output:

 is a modified version of the function that returns Mega bytes instead.

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

Answer: The function takes a number of bytes as input and returns the number of

Bagus, tidak ada waktu tambahan yang dihabiskan untuk menghitung ulang key dan values yang sama untuk lapisan attention! Namun, ada satu hal yang perlu diperhatikan. Meskipun memori puncak yang dibutuhkan untuk $\mathbf{QK}^T$ matrix berkurang secara signifikan, menyimpan key-value cache di memori bisa menjadi sangat mahal dari segi memori untuk urutan input yang panjang atau percakapan multi-turn. Ingatlah bahwa key-value cache perlu menyimpan vektor key-value untuk semua vektor input sebelumnya $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ di semua lapisan self-attention dan untuk semua attention heads.

Mari kita hitung jumlah nilai float yang perlu disimpan dalam cache key-value untuk LLM bigcode/octocoder yang kita gunakan sebelumnya. Jumlah nilai float adalah dua kali panjang sekuens dikalikan jumlah kepala atensi dikalikan dimensi kepala atensi dan dikalikan jumlah lapisan. Menghitung ini untuk LLM kita pada panjang sekuens input hipotetis 16000 menghasilkan:

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

Output:

7864320000

Sekitar 8 miliar nilai float! Menyimpan 8 miliar nilai float dalam presisi float16 membutuhkan sekitar 15 GB RAM, yang kira-kira setengah dari ukuran bobot model itu sendiri! Para peneliti telah mengusulkan dua metode yang memungkinkan pengurangan biaya memori penyimpanan key-value cache secara signifikan, yang dibahas dalam subbagian berikutnya.

3.2.2 Multi-Query-Attention (MQA)

Multi-Query-Attention diusulkan dalam makalah Fast Transformer Decoding: One Write-Head is All You Need karya Noam Shazeer. Sesuai dengan judulnya, Noam menemukan bahwa alih-alih menggunakan bobot proyeksi key-value sebanyak n_head, kita dapat menggunakan satu pasangan bobot proyeksi head-value yang dibagikan di antara semua attention head tanpa menyebabkan penurunan kinerja model secara signifikan.

Dengan menggunakan satu pasangan bobot proyeksi head-value, vektor key value $\mathbf{k}_i, \mathbf{v}_i$ harus identik di semua attention head, yang berarti kita hanya perlu menyimpan 1 pasangan proyeksi key-value di cache alih-alih n_head pasangan.

Karena sebagian besar LLM menggunakan antara 20 hingga 100 kepala perhatian, MQA secara signifikan mengurangi konsumsi memori dari cache kunci-nilai. Untuk LLM yang digunakan dalam notebook ini, kami dapat mengurangi konsumsi memori yang diperlukan dari 15 GB menjadi kurang dari 400 MB pada panjang urutan input 16000.

Selain penghematan memori, MQA juga meningkatkan efisiensi komputasi seperti dijelaskan di bawah ini. Pada decoding auto-regresif, vektor kunci-nilai yang besar perlu dimuat ulang, digabungkan dengan pasangan vektor kunci-nilai saat ini, lalu dimasukkan ke dalam $\mathbf{q}_c\mathbf{K}^T$ perhitungan pada setiap langkah. Untuk decoding auto-regresif, bandwidth memori yang dibutuhkan untuk pemuatan ulang konstan dapat menjadi bottleneck waktu yang serius. Dengan mengurangi ukuran vektor kunci-nilai, jumlah memori yang perlu diakses berkurang, sehingga bottleneck bandwidth memori juga berkurang. Untuk detail lebih lanjut, silakan lihat makalah Noam.

Bagian penting yang perlu dipahami di sini adalah bahwa mengurangi jumlah kepala perhatian key-value menjadi 1 hanya masuk akal jika cache key-value digunakan. Konsumsi memori puncak model untuk satu forward pass tanpa cache key-value tetap tidak berubah karena setiap kepala perhatian masih memiliki vektor query yang unik, sehingga setiap kepala perhatian masih memiliki $\mathbf{QK}^T$ matriks.

MQA telah diadopsi secara luas oleh komunitas dan kini digunakan oleh banyak LLM paling populer:

Selain itu, checkpoint yang digunakan dalam notebook ini - bigcode/octocoder - memanfaatkan MQA.

3.2.3 Grouped-Query-Attention (GQA)

Grouped-Query-Attention, yang diusulkan oleh Ainslie dkk. dari Google, menemukan bahwa penggunaan MQA sering kali dapat menyebabkan penurunan kualitas dibandingkan dengan penggunaan proyeksi multi-key-value head vanilla. Kertas kerja tersebut berargumen bahwa lebih banyak performa model dapat dipertahankan dengan mengurangi jumlah bobot proyeksi query head secara tidak terlalu drastis. Alih-alih menggunakan hanya satu bobot proyeksi key-value, n < n_head bobot proyeksi key-value seharusnya digunakan. Dengan memilih n ke nilai yang jauh lebih kecil daripada n_head, seperti 2, 4, atau 8, hampir semua keuntungan memori dan kecepatan dari MQA dapat dipertahankan sambil mengorbankan lebih sedikit kapasitas model dan dengan demikian, dapat dikatakan, lebih sedikit performa.

Selain itu, para penulis GQA menemukan bahwa checkpoint model yang ada dapat di-uptrain untuk memiliki arsitektur GQA dengan hanya 5% dari komputasi pra-pelatihan aslinya. Meskipun 5% dari komputasi pra-pelatihan asli masih bisa menjadi jumlah yang sangat besar, uptraining GQA memungkinkan checkpoint yang ada tetap berguna untuk urutan input yang lebih panjang.

GQA baru saja diusulkan, sehingga adopsinya masih sedikit pada saat notebook ini ditulis. Aplikasi GQA yang paling menonjol adalah Llama-v2.

Sebagai kesimpulan, sangat disarankan untuk memanfaatkan GQA atau MQA jika LLM di-deploy dengan decoding auto-regresif dan dituntut untuk menangani sekuens input yang besar, seperti yang terjadi misalnya pada chat.

Kesimpulan

Komunitas riset terus-menerus menemukan cara-cara baru yang cerdas untuk mempercepat waktu inferensi pada LLM yang semakin besar. Sebagai contoh, salah satu arah riset yang menjanjikan adalah speculative decoding, di mana “token-token mudah” dihasilkan oleh model bahasa yang lebih kecil dan lebih cepat, dan hanya “token-token sulit” yang dihasilkan oleh LLM itu sendiri. Penjelasan lebih detail berada di luar cakupan notebook ini, tetapi dapat dibaca dalam artikel blog yang bagus ini.

Alasan LLM berskala besar seperti GPT3/4, Llama-2-70b, Claude, dan PaLM dapat berjalan sangat cepat di antarmuka obrolan seperti Hugging Face Chat atau ChatGPT sebagian besar berkat peningkatan presisi, algoritma, dan arsitektur yang disebutkan di atas. Ke depannya, akselerator seperti GPU, TPU, dan lain-lain akan semakin cepat dan memungkinkan memori yang lebih besar, namun kita tetap harus selalu memastikan penggunaan algoritma dan arsitektur terbaik yang tersedia untuk mendapatkan hasil maksimal dari investasi Anda 🤗

Update on GitHub