QLoRA: fine-tuning dużego modelu w 4 bitach

Wyobraź sobie na chwilę, że bierzesz się wreszcie za to, co wiecznie odkładasz „na kiedyś": dostroić model pod siebie — pod swój żargon, swoje dane, swój ton. Wchodzisz na jakiś poradnik o fine-tuningu, czytasz do połowy z zapałem, a potem pada zdanie, które gasi światło: „potrzebujesz kilku kart A100 po 80 GB". No i tyle. Zamykasz kartę w przeglądarce, wracasz do promptowania gotowca i mówisz sobie, że przecież i tak nie miałeś na to czasu. Na forach o fine-tuningu to gatunek posta sam w sobie — ktoś pyta, jak dotrenować siódemkę na własnym korpusie, i pierwsza odpowiedź brzmi „wynajmij GPU w chmurze". Fine-tuning latami był tym pokojem za szybą: widzisz, co w środku, ale klamka jest po drugiej stronie.

W maju 2023 czterech badaczy z University of Washington tę szybę stłukło. Praca nazywa się QLoRA: Efficient Finetuning of Quantized LLMs, a jej główna teza brzmi jak przechwałka z forum: da się dostroić model 65B na jednej karcie 48 GB i zachować jakość pełnego, 16-bitowego fine-tuningu.¹ Nie „prawie", nie „z gwiazdką na jakości". Tak blisko, że w benchmarku różnicy praktycznie nie widać. Jak? Rozłóżmy to na części, bo cały trik jest piękny i składa się z kilku osobnych sprytów, które dopiero razem zaskakują.

Skąd się bierze ten cały VRAM

Zanim rozbijemy QLoRA na części, trzeba wiedzieć, na co przy treningu schodzi pamięć — bo to nie tylko same wagi. Klasyczny fine-tuning musi trzymać na karcie cztery rzeczy naraz: wagi modelu, gradienty (po jednej liczbie na każdą wagę), stany optymalizatora (Adam trzyma ich dwa na wagę) i aktywacje. Przy pełnym treningu w 16 bitach to znaczy, że sama arytmetyka optymalizatora potrafi zająć więcej niż sam model. Dlatego dotrenowanie „całego" modelu 65B to zabawa nie na jedną kartę, a na małą serwerownię.

Pierwszy pomysł, żeby to obejść, jest starszy od QLoRA i nazywa się LoRA (Low-Rank Adapters). Zamiast ruszać wszystkie miliardy wag, zamrażasz je w całości, a obok doklejasz malutkie, niskorangowe macierze — adaptery — i trenujesz tylko je. Nagle gradienty i optymalizator dotyczą ułamka parametrów. To już był przełom. Ale został jeden gruby kamień w bucie: zamrożone wagi bazowe i tak muszą siedzieć na karcie, żeby przez nie puścić forward i backward. A dla modelu 65B te zamrożone wagi w 16 bitach to wciąż jakieś 130 GB. Adaptery są tanie, ale bilet wstępu — sam model — dalej jest za drogi.

Trzy sprytne rzeczy, które robi QLoRA

I tu wchodzi sedno pracy. Skoro te wagi bazowe są zamrożone — nikt ich w treningu nie zmienia — to po co trzymać je w pełnej precyzji? QLoRA mówi: skwantyzujmy je do 4 bitów, zamroźmy je w tej postaci, a gradient przepuśćmy wstecz przez ten zamrożony 4-bitowy model prosto do adapterów LoRA.¹ Trenujemy dalej tylko adaptery — ale bilet wstępu właśnie schudł czterokrotnie. Żeby to zagrało bez utraty jakości, autorzy dorzucili trzy osobne patenty.²

Raz — NF4, czyli 4-bit NormalFloat. Nie wszystkie 4-bitowe typy danych są sobie równe. Wagi wytrenowanej sieci nie są rozrzucone byle jak — układają się mniej więcej w rozkład normalny, taką dzwonową górkę wokół zera. NF4 to typ danych skrojony dokładnie pod ten kształt, teoretycznie optymalny informacyjnie dla wag o rozkładzie normalnym. To ważne: NF4 nie jest zwykłym zmiennoprzecinkowym FP4 z IEEE. To specjalny format, który „marnuje" mniej z tych czterech bitów, bo wie z góry, gdzie wagi się tłoczą.²

Dwa — double quantization (podwójna kwantyzacja). Kwantyzacja blokowa musi gdzieś zapisać stałe skalujące — po jednej na blok wag. To drobne liczby, ale jak masz ich miliony, robi się z tego zauważalny narzut. QLoRA kwantyzuje więc… same stałe kwantyzacji. Kwantyzacja kwantyzacji brzmi jak żart, a oszczędza realnie około 0,4 bita na parametr.² Przy modelu 65B te ułamki bitów zamieniają się w gigabajty.

Trzy — paged optimizers (optymalizatory stronicowane). Przy gradient checkpointingu zużycie pamięci potrafi nagle skoczyć — pojawia się pik i karta krztusi się „out of memory" w najmniej spodziewanym momencie. Paged optimizers używają NVIDIA unified memory, żeby te skoki obsłużyć: gdy na GPU robi się ciasno, stany optymalizatora są po cichu przerzucane do RAM-u i z powrotem, jak stronicowanie pamięci w systemie operacyjnym.¹ Dzięki temu pik nie wywala całego treningu.

To nie jest to samo co GPTQ czy AWQ

Tu łatwo o zamieszanie, więc warto postawić słup graniczny. Kiedy słyszysz „kwantyzacja", najczęściej chodzi o coś zupełnie innego niż QLoRA: o GPTQ albo AWQ, czyli metody, które biorą gotowy, wytrenowany model i ściskają go pod samą inferencję (post-training quantization). Rozgryzaliśmy ten drugi świat osobno tutaj — i to jest inna bajka.

QLoRA nie służy do odchudzania modelu pod odpalanie. QLoRA to metoda fine-tuningu — kwantyzacja 4-bitowa jest w niej tylko trikiem, który pozwala zmieścić trening na jednej karcie. Efekt końcowy to wytrenowane adaptery LoRA, które doklejasz do modelu. Zapamiętaj prostą regułę: GPTQ/AWQ to „jak taniej uruchomić cudzy model", QLoRA to „jak tanio dostroić model pod siebie". Mylenie tych dwóch to klasyk początkującego na forach.

Guanaco, czyli dowód, że to działa

Teoria teorią, ale autorzy poszli na całość: tą metodą dostroili ponad 1000 modeli — różne architektury (LLaMA, T5), różne skale (7B, 13B, 33B, 65B), osiem zestawów danych instrukcyjnych. Skale, których zwykłym fine-tuningiem nikt by na tym sprzęcie nie ruszył.¹

Najlepsza rodzina dostała nazwę Guanaco (tak, po tej andyjskiej kuzynce lamy — w tym świecie wszystko jest lamą). Trenowana na danych OpenAssistant, w benchmarku Vicuna przebiła wszystkie wcześniej publicznie wydane modele i osiągnęła 99,3% poziomu ChatGPT — przy zaledwie 24 godzinach fine-tuningu na jednej karcie.¹

Tu tylko jedna uczciwa gwiazdka, którą stawiają sami autorzy: te 99,3% to wynik z konkretnego benchmarku Vicuna, ocenianego przez GPT-4 — nie jakaś uniwersalna miara „model jest w 99% tak dobry jak ChatGPT do wszystkiego". Praca wprost pokazuje, że ocena GPT-4 to tania i sensowna alternatywa dla oceny ludzkiej, ale zastrzega, że ówczesne benchmarki chatbotów nie były w pełni wiarygodne do precyzyjnego rankingowania modeli.¹ Liczba robi wrażenie — traktuj ją jednak jako „zaskakująco blisko", a nie jako wyrok.

Zrób to sam (a przynajmniej wiedz, gdzie klikać)

Najlepsze jest to, że QLoRA nie została na papierze jako ciekawostka. Kod referencyjny i modele Guanaco wylądowały na licencji MIT w repo artidoro/qlora, razem z jądrami CUDA do treningu 4-bitowego.³ A w praktyce prawie nikt nie pisze tego od zera — całość jest wpięta w ekosystem Hugging Face. Kwantyzację robi biblioteka bitsandbytes, adaptery dokłada PEFT, a spina to wszystko transformers.

Sercem konfiguracji jest BitsAndBytesConfig — i to w nim odnajdziesz jak żywe wszystkie trzy patenty z papieru:⁴

from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                    # wagi bazowe w 4 bitach
    bnb_4bit_quant_type="nf4",            # patent nr 1: NF4, nie zwykły FP4
    bnb_4bit_use_double_quant=True,       # patent nr 2: podwójna kwantyzacja
    bnb_4bit_compute_dtype=torch.bfloat16 # obliczenia w bf16, nie w 4 bitach!
)

Popatrz na ostatnią linijkę — compute_dtype to bf16 (a czym bf16 różni się od f16, rozbieramy osobno), dokładnie ta wyższa precyzja obliczeń, o której gadaliśmy wyżej. Wagi leżą w 4 bitach, ale liczy się porządnie.

A ile realnie karty to zjada? Zależy od modelu, długości sekwencji, batcha i checkpointingu — liczby 65B/48 GB z papieru to konkretny eksperyment, nie uniwersalna tabelka. Ale skala robi wrażenie z drugiej strony też: dzięki nested quantization dokumentacja Hugging Face pokazuje dostrojenie Llama-13B na jednej karcie T4 16 GB — sekwencja 1024 tokeny, batch 1, gradient accumulation na 4 kroki.⁴ T4 to karta, którą dostajesz za darmo w Google Colab.

Jeśli martwi cię, czy w ogóle zmieścisz coś na swojej karcie, to osobna rachunkowość — zbieraliśmy wzory na VRAM tutaj. Ale morał QLoRA jest prosty: bilet wstępu spadł tak, że fine-tuning przestał być salą VIP.

Przez lata fine-tuning działał jak wędkarstwo z burty jachtu — niby każdy może, tylko najpierw kup jacht. QLoRA wsadziła ci w rękę wędkę i pokazała, że ryba bierze tak samo z pomostu za pięć złotych. Zamrożone wagi ściśnięte do czterech bitów, garść adapterów do dotrenowania, jedna karta pod biurkiem — i model, który jeszcze wczoraj wymagał serwerowni, uczy się twojego żargonu przez noc. Nie dlatego, że ktoś dołożył ci pamięci. Dlatego, że ktoś w końcu zapytał, czy naprawdę trzeba trzymać w pełnej precyzji coś, czego i tak nikt w treningu nie rusza.