Zrozumienie parametrów i wymagań pamięciowych dużych modeli językowych: głębokie zanurzenie

Obliczanie liczby parametrów

Pretrained Models For Efficient Transformer Training

Obliczanie parametrów w modelach opartych na Transformatore

Rozbijmy obliczenie parametrów dla każdego składnika modelu opartego na Transformatore. Będziemy używać notacji z oryginalnego artykułu, gdzie d_model reprezentuje wymiar ukrytych stanów modelu.

1. Warstwa wejściowa:
   • Parametry = vocab_size * d_model
2. Uwaga wielogłowa:
   • Dla h głów, z d_k = d_v = d_model / h:
   • Parametry = 4 * d_model^2 (dla Q, K, V i projekcji wyjściowych)
3. Sieć neuronowa feed-forward:
   • Parametry = 2 * d_model * d_ff + d_model + d_ff
   • Gdzie d_ff jest zwykle 4 * d_model
4. Normalizacja warstwy:
   • Parametry = 2 * d_model (dla skali i biasu)

Całkowita liczba parametrów dla jednej warstwy Transformatora:

• Parametry_warstwa = Parametry_uwagi + Parametry_ffn + 2 * Parametry_layernorm

Dla modelu z N warstwami:

• Całkowita liczba parametrów = N * Parametry_warstwa + Parametry_wejściowe + Parametry_wyjściowe

Przykład obliczenia

Rozważmy model o następujących specyfikacjach:

• d_model = 768
• h (liczba głów uwagi) = 12
• N (liczba warstw) = 12
• vocab_size = 50,000

1. Warstwa wejściowa:
   • 50,000 * 768 = 38,400,000
2. Uwaga wielogłowa:
   • 4 * 768^2 = 2,359,296
3. Sieć neuronowa feed-forward:
   • 2 * 768 * (4 * 768) + 768 + (4 * 768) = 4,719,616
4. Normalizacja warstwy:
   • 2 * 768 = 1,536

Całkowita liczba parametrów na warstwę:

• 2,359,296 + 4,719,616 + (2 * 1,536) = 7,081,984

Całkowita liczba parametrów dla 12 warstw:

• 12 * 7,081,984 = 84,983,808

Całkowita liczba parametrów modelu:

• 84,983,808 + 38,400,000 = 123,383,808

Ten model miałby około 123 milionów parametrów.

Rodzaje zużycia pamięci

Podczas pracy z LLM musimy rozważyć dwa główne rodzaje zużycia pamięci:

1. Pamięć modelu: Pamięć wymagana do przechowywania parametrów modelu.
2. Pamięć robocza: Pamięć potrzebna podczas inferencji lub szkolenia do przechowywania pośrednich aktywacji, gradientów i stanów optymalizatora.

Obliczanie pamięci modelu

Pamięć modelu jest bezpośrednio związana z liczbą parametrów. Każdy parametr jest zwykle przechowywany jako 32-bitowa liczba zmiennoprzecinkowa, chociaż niektóre modele używają szkolenia o zmieszanej precyzji z 16-bitowymi liczbami zmiennoprzecinkowymi.

Pamięć modelu (bajty) = Liczba parametrów * Bajty na parametr

Dla naszego przykładowego modelu z 123 milionami parametrów:

• Pamięć modelu (32-bit) = 123,383,808 * 4 bajty = 493,535,232 bajty ≈ 494 MB
• Pamięć modelu (16-bit) = 123,383,808 * 2 bajty = 246,767,616 bajtów ≈ 247 MB

Szacowanie pamięci roboczej

Wymagania pamięci roboczej mogą się znacznie różnić w zależności od konkretnego zadania, rozmiaru partii i długości sekwencji. Szacunkowa pamięć robocza podczas inferencji to:

Pamięć robocza ≈ 2 * Pamięć modelu

To uwzględnia przechowywanie parametrów modelu i pośrednich aktywacji. Podczas szkolenia wymagania pamięci mogą być jeszcze wyższe ze względu na potrzebę przechowywania gradientów i stanów optymalizatora:

Pamięć szkoleniowa ≈ 4 * Pamięć modelu

Dla naszego przykładowego modelu:

• Pamięć robocza podczas inferencji ≈ 2 * 494 MB = 988 MB ≈ 1 GB
• Pamięć szkoleniowa ≈ 4 * 494 MB = 1,976 MB ≈ 2 GB

Stała pamięć i szczytowa pamięć

Podczas szkolenia dużych modeli językowych opartych na architekturze Transformatora, zrozumienie zużycia pamięci jest kluczowe dla efektywnej alokacji zasobów. Rozbijmy wymagania pamięci na dwie główne kategorie: stałą pamięć i szczytową pamięć.

Stała pamięć

Stała pamięć składa się z następujących składników:

1. Wagi modelu: Kopie FP32 parametrów modelu, wymagające 4N bajtów, gdzie N jest liczbą parametrów.
2. Stany optymalizatora: Dla optymalizatora Adam, wymaga to 8N bajtów (2 stany na parametr).
3. Gradients: Kopie FP32 gradientów, wymagające 4N bajtów.
4. Dane wejściowe: Zakładając dane wejściowe int64, wymaga to 8BD bajtów, gdzie B jest rozmiarem partii, a D jest wymiarem danych wejściowych.

Całkowita stała pamięć może być przybliżona przez:

• M_steady = 16N + 8BD bajty

Szczytowa pamięć

Szczytowa pamięć występuje podczas fazy wstecznej, gdy aktywacje są przechowywane do obliczania gradientów. Głównymi składnikami szczytowej pamięci są:

1. Normalizacja warstwy: Wymaga 4E bajtów na warstwę normalizacji, gdzie E = BSH (B: rozmiar partii, S: długość sekwencji, H: rozmiar ukryty).
2. Blok uwagi:
   • Obliczanie QKV: 2E bajty
   • Macierz uwagi: 4BSS bajty (S: długość sekwencji)
   • Wyjście uwagi: 2E bajty
3. Blok feed-forward:
   • Pierwsza warstwa liniowa: 2E bajty
   • Aktywacja GELU: 8E bajtów
   • Druga warstwa liniowa: 2E bajty
4. Strata cross-entropy:
   • Logity: 6BSV bajtów (V: rozmiar słownika)

Całkowita pamięć aktywacji może być oszacowana jako:

• M_act = L * (14E + 4BSS) + 6BSV bajtów

Gdzie L jest liczbą warstw Transformatora.

Całkowita szczytowa pamięć

Całkowita szczytowa pamięć podczas szkolenia może być przybliżona przez połączenie stałej pamięci i pamięci aktywacji:

• M_peak = M_steady + M_act + 4BSV bajty

Dodatkowy termin 4BSV uwzględnia dodatkową alokację na początku fazy wstecznej.

Poprzez zrozumienie tych składników, możemy zoptymalizować zużycie pamięci podczas szkolenia i inferencji, zapewniając efektywną alokację zasobów i poprawiając wydajność dużych modeli językowych.