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

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

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 reguła jest:

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

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

Podczas inferencji, ślad pamięciowy jest dominowany przez parametry modelu i tymczasowe tensory aktywacyjne. Szacunek wysoko poziomowy dla maksymalnego zużycia pamięci podczas inferencji jest sumą pamięci wymaganej do załadowania parametrów modelu i pamięci dla aktywacji.

Ilościowy szacunek pamięci inferencyjnej

Rozważmy szacunek pamięci dla inferencji przy użyciu 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

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Wynik:

Maksymalne 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 inferencyjnej z kwantyzacją

Podczas gdy bfloat16 jest powszechnie używana precyzja dla treningu LLM, badacze odkryli, że kwantyzacja wag modelu do niższej precyzji, takiej 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ń inferencyjnych, takich jak generowanie tekstu.

Zobaczmy oszczędności pamięci z kwantyzacji 8-bitowej i 4-bitowej 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 zmniejsza ją jeszcze bardziej do zaledwie 9,5 GB! To pozwala na uruchomienie modelu OctoCode z 15 miliardami 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 precyzją 8-bitową lub bfloat16. 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 na środowiskach o ograniczonych zasobach, takich jak instancje chmury, urządzenia brzegowe lub nawet telefony komórkowe, znacznie zmniejszając ślad pamięciowy.

Szacowanie pamięci dla dostrajania

Podczas gdy kwantyzacja jest głównie używana do wydajnej inferencji, techniki takie jak tensor parallelism i model parallelism są niezbędne do zarządzania wymaganiami pamięciowymi podczas treningu lub dostrajania dużych modeli językowych.

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

• Gradientów
• Stanów optymalizatora
• Aktywacji z przodu przechowywanych do wstecznej propagacji

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 z 7 miliardami 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 stosowania technik dostrajania rozproszonego.

Techniki dostrajania rozproszonego

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

1. Parallelizm danych: Klasyczne podejście parallelizmu danych replikuje cały model na wielu GPU, podczas gdy dzieli i rozdziela partie danych szkoleniowych. To zmniejsza czas treningu liniowo z liczbą GPU, ale nie zmniejsza maksymalnego zużycia pamięci na każdym GPU.
2. ZeRO Stage 3: Zaawansowana forma parallelizmu danych, która partitionuje parametry modelu, gradienty i stany optymalizatora na GPU. To zmniejsza pamięć w porównaniu z klasycznym parallelizmem danych, przechowując tylko wymagane dane partitionowane na każdym GPU podczas różnych faz treningu.
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 partitionowanym zestawie parametrów, gradientów i stanów optymalizatora, co prowadzi do znacznych oszczędności pamięci.
4. Parallelizm potokowy: Ta technika partitionuje warstwy modelu na różne GPU/roboty, z których każdy wykonuje podzbiór warstw. Aktywacje są przekazywane między robotami, co zmniejsza maksymalne zużycie pamięci, ale zwiększa nakład komunikacyjny.

Szacowanie zużycia pamięci dla tych metod rozproszonych jest niełatwe, ponieważ dystrybucja parametrów, gradientów, aktywacji i stanów optymalizatora 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 przydziału 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 przydziału pamięci dla ciała transformatora i głowy modelu językowego.

Wyniki eksperymentalne pokazują, że LLMem może szacować maksymalne 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%.