Ile kosztuje trening LLM: 6ND FLOP-ów, 16 bajtów na parametr, Chinchilla
Rachunek kosztu treningu na serwetce, któremu można ufać: reguła 6ND dla obliczeń, co Chinchilla mówi o liczbie tokenów na parametr, dlaczego Adam kosztuje 16 bajtów na parametr i na co idzie pamięć aktywacji. Policzone dla modelu 125M, modelu 1B i fine-tuningu modelu 4,5B na jednej RTX 4090.
Zanim zaczniesz cokolwiek trenować, możesz z grubsza przewidzieć, ile to będzie kosztować w godzinach i w pamięci, na podstawie trzech liczb: parametrów, tokenów i wydajności GPU. Poniżej rachunek, którego używam, i skąd bierze się każda z tych liczb.
Obliczenia: 6ND
Trening gęstego transformera (dense) o N parametrach na D tokenach kosztuje około:
C ≈ 6 · N · D FLOPsPrzebieg w przód to około 2N FLOP-ów na token, po jednym mnożeniu i jednym dodawaniu na wagę. Przebieg wsteczny to mniej więcej dwa razy tyle, 4N, bo liczy gradienty zarówno dla aktywacji, jak i dla wag. Atencja dodaje składnik, który rośnie z długością kontekstu, ale przy typowych długościach treningowych, dla modeli tej wielkości, to niewielka poprawka (Kaplan i in., 2020).
Żeby zamienić FLOP-y na godziny, potrzebujesz wydajności, którą naprawdę uzyskasz. Szczytowej wydajności rdzeni tensorowych nie osiąga się nigdy. Wykorzystanie FLOP-ów modelu (MFU) na poziomie 30–50% jest typowe dla dobrze dostrojonego treningu na jednym GPU. Dla RTX 4090 przyjmij około 165 TFLOPS szczytowo w bf16 z akumulacją w fp32 i 40% wykorzystania, czyli 66 TFLOPS wydajności ciągłej:
| Trening | N | D | FLOP-y | Czas przy 66 TFLOPS |
|---|---|---|---|---|
| Mały model, optymalny obliczeniowo | 125M | 2,5 mld tokenów | 1,9 × 10¹⁸ | ~7,9 godziny |
| Model 1B, optymalny obliczeniowo | 1B | 20 mld tokenów | 1,2 × 10²⁰ | ~21 dni |
| Fine-tuning 4,5B | 4,5B | 1 mld tokenów | 2,7 × 10¹⁹ | ~4,7 dnia |
Pierwszy wiersz to model takiego rodzaju, jaki od zera trenuje moje repozytorium llm-mini, żeby pokazać, jak poszczególne elementy do siebie pasują. Ostatni wiersz pokazuje, jak wygląda fine-tuning modelu takiego jak compass.
Ile tokenów: Chinchilla
Hoffmann i in. (2022) ustalili, że przy stałym budżecie obliczeniowym model optymalny pod względem straty trenuje się na około 20 tokenach na parametr: 2,5 mld tokenów dla modelu 125M, 20 mld dla modelu 1B. Według tej miary wcześniejsze duże modele były mocno niedotrenowane.
W praktyce liczą się dwa zastrzeżenia. Chinchilla optymalizuje obliczenia treningowe, a nie wdrożenie. Jeśli model ma serwować miliardy tokenów, opłaca się trenować mniejszy model znacznie dłużej niż 20 tokenów na parametr, bo za inferencję płaci się przy każdym użyciu. Większość najnowszych otwartych modeli robi dokładnie to, przy setkach albo tysiącach tokenów na parametr. A budżety fine-tuningu wyznacza zmiana zachowania, którą chcesz uzyskać, a nie Chinchilla.
Pamięć: 16 bajtów na parametr
Trening w mieszanej precyzji z Adamem przechowuje na każdy parametr:
wagi bf16 2 bajty
gradienty bf16 2
wagi główne fp32 4
pierwszy moment Adama 4
drugi moment Adama 4
── 16 bajtów na parametr (rachunek z pracy o ZeRO)Dla modelu 4,5B to 72 GB jeszcze przed pierwszą aktywacją, trzy razy więcej, niż mieści jedna karta 24 GB. Możliwe wyjścia:
- 8-bitowe stany optymalizatora. Trzymanie momentów Adama w 8 bitach sprowadza to do około 10 bajtów na parametr, tutaj 45 GB. Wciąż za dużo.
- LoRA albo QLoRA. Trenuj tylko adaptery, a 16 bajtów modelu bazowego spada do 2, albo do około 0,5 przy 4 bitach. Liczby są w notatce o LoRA.
- Offload albo shardowanie. Przenieś stany optymalizatora do RAM-u CPU (ZeRO-Offload) albo rozłóż je na wiele GPU. To wymiana pamięci na przepustowość.
Aktywacje: część, która rośnie z kontekstem
Aktywacje zależą od rozmiaru partii b, długości sekwencji s, wymiaru ukrytego h i liczby głowic atencji a. Korthikanti i in. (2022) podają koszt na warstwę dla bloku w stylu GPT, bez ponownego obliczania, jako około s·b·h·(34 + 5·a·s/h) bajtów.
Weźmy konfigurację klasy 4,5B (h = 3072, a = 24) przy s = 4096, b = 1. To około 2,4 GB na warstwę, z czego prawie wszystko przypada na składnik wyników atencji 5·a·s/h. FlashAttention nigdy nie zapisuje wyników atencji do pamięci, więc ten składnik w zasadzie znika i zostaje około 0,43 GB na warstwę. Gradient checkpointing obcina większość reszty, licząc aktywacje ponownie w przebiegu wstecznym. Kosztuje to mniej więcej o jedną trzecią więcej obliczeń, ale zwykle to właśnie ta wymiana sprawia, że trening się mieści.
Zasady, których się trzymam
- Szacuj, zanim zaczniesz:
6NDprzy realistycznym wykorzystaniu 30–50% i 16 bajtów na trenowany parametr. - Trenuj małe modele długo, jeśli mają być serwowane. Chinchilla dotyczy kosztu treningu, a nie kosztu całkowitego.
- Żeby zmieścić się w pamięci, sięgaj najpierw po adaptery, potem po 8-bitowe optymalizatory, potem po checkpointing, a na końcu po offload, w tej kolejności, od najmniej bolesnego.
- Mierz wykorzystanie wcześnie. Jeśli krótki przebieg pokaże 15% MFU, popraw ładowanie danych albo jądra, zanim zdecydujesz się na wiele dni treningu.
Źródła: Kaplan i in., Scaling Laws for Neural Language Models (2020). Hoffmann i in., Training Compute-Optimal Large Language Models (2022). Rajbhandari i in., ZeRO (2020). Korthikanti i in., Reducing Activation Recomputation in Large Transformer Models (2022). Czasy to szacunki przy ciągłej wydajności 66 TFLOPS, a nie pomiary.