Matematyka LoRA i QLoRA: rząd, alfa, pamięć i co naprawdę się trenuje

Fine-tuning wszystkich wag modelu 7B wymaga ponad 100 GB pamięci GPU. LoRA trenuje zamiast tego wyraźnie poniżej 1% parametrów i wypada zaskakująco blisko. QLoRA idzie dalej i trzyma zamrożony model w 4 bitach. Poniżej arytmetyka obu metod i to, gdzie za ten skrót płacisz.

Pomysł

Dostrojona macierz wag to oryginał plus zmiana: W' = W + ΔW. LoRA (Hu i in., 2021) stawia na to, że ta zmiana ma niski rząd (rank), i trenuje ją jako iloczyn dwóch wąskich macierzy:

W' = W + (α / r) · B · A
W: d × k (zamrożona)    B: d × r    A: r × k    r ≪ min(d, k)

A jest inicjalizowana losowo, a B zerami, więc model zaczyna dokładnie tam, gdzie był. Gradienty dostają tylko A i B.

Ile to parametrów

Jedna projekcja 4096 × 4096 ma 16,8 mln wag. Adapter rzędu 16 na tej projekcji ma 16 × (4096 + 4096) = 131 072 parametry, czyli 0,78% macierzy.

W całym modelu o wymiarach Llama-2-7B (32 warstwy, wymiar ukryty 4096, szerokość MLP 11008), z adapterami rzędu 16 na wszystkich siedmiu warstwach liniowych w każdym bloku:

q, k, v, o:   4 × 16 × (4096 + 4096)  = 524 288
gate, up:     2 × 16 × (4096 + 11008) = 483 328
down:         1 × 16 × (11008 + 4096) = 241 664
na warstwę                            = 1 249 280
× 32 warstwy                          ≈ 40,0 mln   (0,59% z 6,74 mld)

Czterdzieści milionów trenowalnych parametrów z 6,74 miliarda.

Dlaczego zużycie pamięci gwałtownie spada

Pełny fine-tuning z Adamem w mieszanej precyzji kosztuje około 16 bajtów na parametr:

• 2 bajty na wagi w bf16 i 2 na gradienty;
• 4 na główną kopię w fp32;
• 4 + 4 na dwa momenty Adama.

Przy 6,74 mld parametrów to około 108 GB, jeszcze bez aktywacji.

Z LoRA zamrożony model bazowy potrzebuje tylko swoich wag: 13,5 GB w bf16. Te 16 bajtów na parametr dotyczy tylko 40 mln parametrów adaptera, czyli około 0,64 GB. Do tego dochodzą aktywacje i właśnie tu gradient checkpointing zarabia na siebie.

QLoRA (Dettmers i in., 2023) trzyma zamrożony model bazowy w 4-bitowym NF4, typie danych z poziomami rozmieszczonymi pod wagi o rozkładzie normalnym. Kwantyzuje też same stałe kwantyzacji (podwójna kwantyzacja), co daje łącznie około 4,127 bita na parametr. Model bazowy jest dekwantyzowany w locie w przebiegu w przód, a gradienty płyną przez niego do adapterów w bf16:

Model bazowy	Zamrożony w bf16	Zamrożony w NF4 (QLoRA)
7B	13,5 GB	3,5 GB
27B	54 GB	13,9 GB

Dlatego model 27B w ogóle da się zaadaptować na jednej karcie 24 GB: 14 GB na model bazowy, małe adaptery, a reszta na aktywacje przy umiarkowanych rozmiarach partii i długościach sekwencji.

Co robi alfa

Aktualizacja jest skalowana przez α / r. Zwiększ rząd przy stałym α, a każdy z kierunków rzędu wnosi mniej, więc efektywny współczynnik uczenia (learning rate) zmienia się wraz z rzędem. Wiele konfiguracji ustawia α na r albo 2r i na tym poprzestaje. rsLoRA (Kalajdzievski, 2023) opowiada się za α / √r, co utrzymuje stabilność aktualizacji przy rosnącym rzędzie. Niezależnie od wariantu, po zmianie rzędu sprawdź ponownie współczynnik uczenia.

Gdzie umieszczać adaptery

Oryginalna praca adaptowała tylko projekcje atencji. Autorzy QLoRA stwierdzili, że do dorównania jakości pełnego fine-tuningu potrzebna była adaptacja każdej warstwy liniowej, łącznie z MLP. To kosztuje więcej parametrów, ale przy tych rozmiarach są tanie. Domyślnie adaptuję wszystkie warstwy liniowe.

W czym LoRA jest dobra, a w czym słaba

LoRA dobrze uczy zachowania: formatu, tonu, dyscypliny wywołań narzędzi, stylu rozumowania, konwencji danej dziedziny. Do tego służy moje repozytorium lora: to mały, skuteczny sposób na fine-tuning modelu rozumującego.

Słabiej radzi sobie z dodawaniem dużych ilości nowej wiedzy. Biderman i in. (2024) stwierdzili, że przy kontynuowanym pre-trainingu na kodzie i matematyce LoRA uczy się mniej niż pełny fine-tuning, a zarazem mniej zapomina z tego, co wiedział model bazowy. Taki jest kompromis. Jeśli zadanie brzmi „poznaj dogłębnie nową dziedzinę”, potrzebujesz pełnego fine-tuningu albo znacznie wyższego rzędu. Jeśli brzmi „zachowuj się tak, i to niezawodnie”, LoRA zwykle bardziej się opłaca.

Zasady, których się trzymam

1. Adaptuj wszystkie warstwy liniowe przy rzędzie 16–64, zanim sięgniesz po wyższy rząd.
2. Traktuj rząd i współczynnik uczenia jako parę. Zmieniasz jedno, stroisz na nowo drugie.
3. Używaj QLoRA, gdy model bazowy nie mieści się w bf16. Baza w NF4 kosztuje trochę jakości; rezygnacja z treningu kosztuje więcej.
4. Scalaj do serwowania. W + (α/r)·BA składa się z powrotem w jedną macierz, więc inferencja nie płaci nic dodatkowo. Potem ponownie skwantyzuj scalone wagi do wdrożenia.

Źródła: Hu i in., LoRA (2021). Dettmers i in., QLoRA (2023). Kalajdzievski, A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA (2023). Biderman i in., LoRA Learns Less and Forgets Less (2024). Liczby parametrów są policzone dla modelu o wymiarach Llama-2-7B.