Optymalizacja pamięci dla wnioskowania i dostrajania dużych modeli językowych

Duże modele językowe (LLM) takie jak GPT-4, Bloom i LLaMA osiągnęły znaczące możliwości dzięki skalowaniu do miliardów parametrów. Jednak wdrożenie tych ogromnych modeli do wnioskowania lub dostrajania jest wyzwaniem ze względu na ich ogromne wymagania pamięciowe. W tym technicznym blogu będziemy eksplorować techniki szacowania i optymalizacji zużycia pamięci podczas wnioskowania i dostrajania LLM na różnych konfiguracjach sprzętowych.

Zrozumienie wymagań pamięciowych

Pamięć wymagana do załadowania LLM jest głównie determinowana przez liczbę parametrów i precyzję numeryczną używaną do przechowywania parametrów. Prosta zasada to:

• Ładowanie modelu z X miliardami parametrów wymaga około 4X GB pamięci VRAM w precyzji 32-bit float
• Ładowanie modelu z X miliardami parametrów wymaga około 2X GB pamięci VRAM w precyzji 16-bit bfloat16/float16

Na przykład, ładowanie modelu GPT-3 o 175 miliardach parametrów wymagałoby około 350 GB pamięci VRAM w precyzji bfloat16. Jak dotąd, największe komercyjnie dostępne karty graficzne, takie jak NVIDIA A100 i H100, oferują tylko 80 GB pamięci VRAM, co wymaga techniki tensorowej i modelowej.

Podczas wnioskowania, ślad pamięciowy jest dominowany przez parametry modelu i tymczasowe tensory aktywacyjne. Szacunek wysokiego poziomu dla szczytowego zużycia pamięci podczas wnioskowania jest suma pamięci wymaganej do załadowania parametrów modelu i pamięci dla aktywacji.

Ilościowy szacunek pamięci wnioskowania

Rozważmy szacunek pamięci wnioskowania za pomocą modelu OctoCode, który ma około 15 miliardów parametrów w formacie bfloat16 (~ 31 GB). Użyjemy biblioteki Transformers do załadowania modelu i wygenerowania tekstu:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Wynik:

Szczytowe zużycie pamięci GPU wynosi około 29 GB, co odpowiada naszemu szacunkowi 31 GB dla załadowania parametrów modelu w formacie bfloat16.

Optymalizacja pamięci wnioskowania z kwantyzacją

Chociaż bfloat16 jest powszechnie używaną precyzją dla szkolenia LLM, badacze odkryli, że kwantyzacja wag modelu do niższych typów danych, takich jak 8-bitowe liczby całkowite (int8) lub 4-bitowe liczby całkowite, może znacznie zmniejszyć zużycie pamięci z minimalną utratą dokładności dla zadań wnioskowania, takich jak generowanie tekstu.

Zobaczmy oszczędności pamięci z kwantyzacją 8-bitową i 4-bitową modelu OctoCode:

</div>
# 8-bitowa kwantyzacja
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=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):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Wynik:

# 4-bitowa kwantyzacja
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=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):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Wynik:

Z kwantyzacją 8-bitową, wymaganie pamięci spada z 31 GB do 15 GB, podczas gdy kwantyzacja 4-bitowa redukuje ją jeszcze bardziej do zaledwie 9,5 GB! To pozwala na uruchomienie modelu OctoCode o 15 miliardach parametrów na konsumentach GPU, takich jak RTX 3090 (24 GB VRAM).

Jednak należy zauważyć, że bardziej agresywna kwantyzacja, taka jak 4-bitowa, może czasem prowadzić do degradacji dokładności w porównaniu z 8-bitową lub bfloat16 precyzją. Istnieje kompromis między oszczędnościami pamięci a dokładnością, który użytkownicy powinni ocenić dla swojego przypadku użycia.

Kwantyzacja jest potężną techniką, która może umożliwić wdrożenie LLM w środowiskach o ograniczonych zasobach, takich jak instancje chmurowe, urządzenia brzegowe lub nawet telefony komórkowe, znacznie redukując ślad pamięciowy.

Szacowanie pamięci dla dostrajania

Podczas gdy kwantyzacja jest głównie używana do efektywnego wnioskowania, techniki takie jak tensorowa i modelowa są niezbędne do zarządzania wymaganiami pamięciowymi podczas szkolenia lub dostrajania dużych modeli językowych.

Szczytowe zużycie pamięci podczas dostrajania jest zwykle 3-4 razy wyższe niż podczas wnioskowania ze względu na dodatkowe wymagania pamięciowe dla:

• Gradientów
• Stanów optymalizacji
• Aktywacji z przodu przechowywanych do wstecznego propagowania

Konserwatywny szacunek jest taki, że dostrajanie LLM z X miliardami parametrów wymaga około 4 * (2X) = 8X GB pamięci VRAM w precyzji bfloat16.

Na przykład, dostrajanie modelu LLaMA o 7 miliardach parametrów wymagałoby około 7 * 8 = 56 GB pamięci VRAM na GPU w precyzji bfloat16. To przekracza pojemność pamięciową obecnych GPU, co wymaga technik dostrajania rozproszonego.

Techniki dostrajania rozproszonego

Kilka metod dostrajania rozproszonego zostało zaproponowanych, aby pokonać ograniczenia pamięciowe GPU dla dużych modeli:

1. Parallelizm danych: Klasyczny parallelizm danych replikuje cały model na wielu GPU, podczas gdy dzieli i rozdziela partie danych szkoleniowych. To redukuje czas szkolenia liniowo z liczbą GPU, ale nie redukuje szczytowego wymagania pamięciowego na każdym GPU.
2. ZeRO Stage 3: Zaawansowana forma parallelizmu danych, która dzieli parametry modelu, gradienty i stany optymalizacji na GPU. To redukuje pamięć w porównaniu z klasycznym parallelizmem danych, przechowując tylko wymagane dane podzielone na każdym GPU podczas różnych faz szkolenia.
3. Parallelizm tensorowy: Zamiast replikować model, parallelizm tensorowy dzieli parametry modelu na wiersze lub kolumny i rozdziela je na GPU. Każdy GPU operuje na podziale parametrów, gradientów i stanów optymalizacji, co prowadzi do znacznych oszczędności pamięci.
4. Parallelizm potokowy: Ta technika dzieli warstwy modelu na różne GPU/roboty, z których każdy wykonuje podzbiór warstw. Aktywacje są przekazywane między robotami, redukując szczytowe zużycie pamięci, ale zwiększając nakład komunikacyjny.

Szacowanie zużycia pamięci dla tych metod rozproszonych jest nietrywialne, ponieważ dystrybucja parametrów, gradientów, aktywacji i stanów optymalizacji różni się w zależności od techniki. Ponadto różne komponenty, takie jak ciało transformatora i głowa modelu językowego, mogą wykazywać różne zachowania alokacji pamięci.

Rozwiązanie LLMem

Naukowcy niedawno zaproponowali LLMem, rozwiązanie, które dokładnie szacuje zużycie pamięci GPU podczas stosowania metod dostrajania rozproszonego do LLM na wielu GPU.

Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLM

LLMem uwzględnia czynniki, takie jak ponowne łączenie parametrów przed obliczeniami (ZeRO Stage 3), gromadzenie danych wyjściowych w fazie wstecznej (parallelizm tensorowy) oraz różne strategie alokacji pamięci dla ciała transformatora i głowy modelu językowego.

Wyniki eksperymentalne pokazują, że LLMem może szacować szczytowe zużycie pamięci GPU dla dostrajania LLM na jednym GPU z błędem do 1,6%, przewyższając średni błąd DNNMem o 42,6%. Podczas stosowania metod dostrajania rozproszonego do LLM z ponad miliardem parametrów na wielu GPU, LLMem osiąga imponujący średni błąd 3,0%.