Optymalizacja pamięci pod kątem wnioskowania i dostrajania modelu wielkojęzykowego

Duże modele językowe (LLM), takie jak GPT-4, Bloom i LLaMA, osiągnęły niezwykłe możliwości poprzez skalowanie do miliardów parametrów. Jednak wdrażanie tych ogromnych modeli w celu wnioskowania lub dostrajania jest wyzwaniem ze względu na ich ogromne wymagania dotyczące pamięci. W tym blogu technicznym omówimy techniki szacowania i optymalizacji zużycia pamięci podczas wnioskowania LLM i dostrajania różnych konfiguracji sprzętowych.

Zrozumienie wymagań dotyczących pamięci

Pamięć wymagana do załadowania LLM zależy przede wszystkim od liczby parametrów i precyzji numerycznej użytej do przechowywania parametrów. Prosta zasada jest następująca:

• Załadowanie modelu z X miliardami parametrów wymaga mniej więcej 4X GB VRAM w 32-bit precyzja pływania
• Załadowanie modelu z X miliardami parametrów wymaga mniej więcej 2X GB VRAM w 16-bit precyzja bfloat16/float16

Na przykład załadowanie modelu GPT-175 z parametrem 3B wymagałoby około 350 GB pamięci VRAM z precyzją bfloat16. Na dzień dzisiejszy największe dostępne na rynku procesory graficzne, takie jak NVIDIA A100 i H100, oferują tylko 80 GB pamięci VRAM, co wymaga stosowania technik równoległości tensorów i równoległości modeli.

Podczas wnioskowania wielkość pamięci jest zdominowana przez parametry modelu i wytworzone tymczasowe tensory aktywacji. Ogólna ocena szczytowego użycia pamięci podczas wnioskowania to suma pamięci wymaganej do załadowania parametrów modelu i pamięci do aktywacji.

Kwantyfikacja pamięci wnioskowania

Określmy ilościowo wymagania dotyczące pamięci do wnioskowania, korzystając z modelu OctoCode, który ma około 15 miliardów parametrów w formacie bfloat16 (~ 31 GB). Skorzystamy z Biblioteka transformatorów aby załadować model i wygenerować tekst:

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())

Wyjście:

Szczytowe wykorzystanie pamięci GPU wynosi około 29 GB, co jest zgodne z naszymi szacunkami na 31 GB przy ładowaniu parametrów modelu w formacie bfloat16.

Optymalizacja pamięci wnioskowania za pomocą kwantyzacji

Chociaż bfloat16 jest powszechną precyzją stosowaną do uczenia LLM, badacze odkryli, że kwantowanie wag modelu do typów danych o niższej precyzji, takich jak 8-bitowe liczby całkowite (int8) lub 4-bitowe liczby całkowite, może znacznie zmniejszyć zużycie pamięci przy minimalnej utracie dokładności w przypadku zadań wnioskowania, takich jak generowanie tekstu.

Zobaczmy oszczędność pamięci dzięki 8-bitowej i 4-bitowej kwantyzacji modelu OctoCode:

</div>
# 8-bit quantization
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>

Wyjście:

# 4-bit quantization
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())

Wyjście:

Przy 8-bitowej kwantyzacji zapotrzebowanie na pamięć spada z 31 GB do 15 GB, podczas gdy 4-bitowa kwantyzacja dodatkowo zmniejsza je do zaledwie 9.5 GB! Umożliwia to uruchomienie modelu OctoCode z parametrem 15B na konsumenckich procesorach graficznych, takich jak RTX 3090 (24 GB VRAM).

Należy jednak pamiętać, że bardziej agresywna kwantyzacja, np. 4-bitowa, może czasami prowadzić do pogorszenia dokładności w porównaniu z precyzją 8-bitową lub bfloat16. Istnieje kompromis między oszczędnością pamięci a dokładnością, który użytkownicy powinni ocenić pod kątem swojego przypadku użycia.

Kwantyzacja to zaawansowana technika, która może umożliwić wdrożenie LLM w środowiskach o ograniczonych zasobach, takich jak instancje w chmurze, urządzenia brzegowe, a nawet telefony komórkowe, drastycznie zmniejszając zużycie pamięci.

Szacowanie pamięci w celu dostrajania

Chociaż kwantyzacja jest wykorzystywana głównie do wydajnego wnioskowania, techniki takie jak równoległość tensorów i równoległość modeli mają kluczowe znaczenie w zarządzaniu wymaganiami dotyczącymi pamięci podczas uczenia lub strojenie dużych modeli językowych.

Szczytowe zużycie pamięci podczas dostrajania jest zwykle 3–4 razy większe niż wnioskowanie ze względu na dodatkowe wymagania dotyczące pamięci dla:

• Gradienty
• Stany optymalizatora
• Aktywacje z przebiegu w przód przechowywane w celu propagacji wstecznej

Ostrożne szacunki mówią, że dostrojenie LLM przy użyciu X miliardów parametrów wymaga ok 4 * (2X) = 8X GB VRAM z precyzją bfloat16.

Na przykład dostrojenie modelu LLaMA z parametrem 7B wymagałoby w przybliżeniu 7*8 = 56 GB pamięci VRAM na procesor graficzny z precyzją bfloat16. Przekracza to pojemność pamięci obecnych procesorów graficznych, co wymaga rozproszonych technik dostrajania.

Rozproszone techniki dostrajania

Zaproponowano kilka rozproszonych metod dostrajania w celu przezwyciężenia ograniczeń pamięci GPU w przypadku dużych modeli:

1. Równoległość danych: Klasyczne podejście do równoległości danych replikuje cały model na wielu procesorach graficznych, jednocześnie dzieląc i dystrybuując partie danych szkoleniowych. Skraca to czas szkolenia liniowo wraz z liczbą procesorów graficznych, ale nie zmniejsza szczytowego zapotrzebowania na pamięć na każdym procesorze graficznym.
2. Etap 3 ZeroRO: Zaawansowana forma równoległości danych, która dzieli parametry modelu, gradienty i stany optymalizatora pomiędzy procesory graficzne. Zmniejsza pamięć w porównaniu z klasyczną równoległością danych, przechowując tylko wymagane, podzielone na partycje dane na każdym procesorze graficznym podczas różnych faz szkolenia.
3. Równoległość tensorowa: Zamiast replikować model, równoległość tensorowa dzieli parametry modelu na wiersze lub kolumny i rozdziela je pomiędzy procesory graficzne. Każdy procesor graficzny działa na podzielonym zestawie parametrów, gradientów i stanów optymalizatora, co prowadzi do znacznych oszczędności pamięci.
4. Równoległość rurociągu: Ta technika dzieli warstwy modelu na różne procesory graficzne/procesy robocze, przy czym każde urządzenie wykonuje podzbiór warstw. Aktywacje są przekazywane między procesami roboczymi, zmniejszając pamięć szczytową, ale zwiększając obciążenie komunikacyjne.

Oszacowanie wykorzystania pamięci dla tych metod rozproszonych nie jest trywialne, ponieważ rozkład parametrów, gradientów, aktywacji i stanów optymalizatora różni się w zależności od techniki. Co więcej, różne komponenty, takie jak korpus transformatora i głowica modelująca język, mogą wykazywać różne zachowania przy alokacji pamięci.

Rozwiązanie LLMem

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

Szacowanie użycia pamięci GPU w celu dostrajania wstępnie przeszkolonego LLM

LLMem bierze pod uwagę takie czynniki, jak rekombinacja parametrów przed obliczeniami (etap 3 ZeRO), gromadzenie danych wyjściowych w przebiegu wstecznym (równoległość tensorów) oraz różne strategie alokacji pamięci dla korpusu transformatora i głowicy modelującej język.

Wyniki eksperymentów pokazują, że LLMem może oszacować szczytowe wykorzystanie pamięci GPU w celu dostrojenia LLM na pojedynczym GPU z poziomem błędów do 1.6%, co przewyższa średni poziom błędów najnowocześniejszego DNNMem wynoszący 42.6%. Stosując rozproszone metody dostrajania do LLM z ponad miliardem parametrów na wielu procesorach graficznych, LLMem osiąga imponujący średni poziom błędów wynoszący 3.0%.