Оптимізація пам’яті для висновку та налаштування великих мовних моделей

Великі мовні моделі (LLM) типу GPT-4, Bloom, і LLaMA досягли вражаючих можливостей завдяки масштабуванню до мільярдів параметрів. Однак розгортання цих масивних моделей для висновку або налаштування є складним завданням через їх величезні вимоги до пам’яті. У цьому технічному блозі ми дослідимо техніки оцінки та оптимізації споживання пам’яті під час висновку та налаштування LLM на різних апаратних конфігураціях.

Поняття вимог до пам’яті

Пам’ять, необхідна для завантаження LLM, визначається головним чином кількістю параметрів і числовою точністю, яку використовують для зберігання параметрів. Просте правило:

• Завантаження моделі з X мільярдами параметрів вимагає приблизно 4X ГБ відеопам’яті у форматі 32-біт float
• Завантаження моделі з X мільярдами параметрів вимагає приблизно 2X ГБ відеопам’яті у форматі 16-біт bfloat16/float16

Наприклад, завантаження моделі GPT-3 з 175 мільярдами параметрів вимагатиме приблизно 350 ГБ відеопам’яті у форматі bfloat16. На сьогодні найбільші комерційно доступні графічні процесори типу NVIDIA A100 і H100 пропонують лише 80 ГБ відеопам’яті, що вимагає технік паралелізму тензорів і паралелізму моделей.

Під час висновку пам’ять домінується параметрами моделі та тимчасовими тензорами активації. Вищий рівень оцінки пікової пам’яті під час висновку – це сума пам’яті, необхідної для завантаження параметрів моделі, і пам’яті для активації.

Кількісна оцінка висновку пам’яті

Давайте кількісно оцінимо вимоги до пам’яті для висновку за допомогою моделі OctoCode, яка має близько 15 мільярдів параметрів у форматі bfloat16 (~ 31 ГБ). Ми використаємо Transformers library для завантаження моделі та генерації тексту:

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

Вихід:

Пікове використання відеопам’яті складає близько 29 ГБ, що відповідає нашій оцінці 31 ГБ для завантаження параметрів моделі у форматі bfloat16.

Оптимізація висновку пам’яті за допомогою квантування

Хоча bfloat16 є звичайною точністю, використовуваною для навчання LLM, дослідники виявили, що квантування ваг моделі до нижчої точності даних, таких як 8-бітові цілі числа (int8) або 4-бітові цілі числа, може суттєво зменшити використання пам’яті з мінімальною втратою точності для завдань висновку, таких як генерація тексту.

Давайте розглянемо економію пам’яті від 8-бітового та 4-бітового квантування моделі OctoCode:

</div>
# 8-бітове квантування
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>

Вихід:

# 4-бітове квантування
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())

Вихід:

З 8-бітовим квантуванням вимоги до пам’яті знижуються з 31 ГБ до 15 ГБ, тоді як 4-бітове квантування ще більше знижує їх до 9,5 ГБ! Це дозволяє запускати модель OctoCode з 15 мільярдами параметрів на споживчій графічній карті типу RTX 3090 (24 ГБ відеопам’яті).

Однак зверніть увагу, що більш агресивне квантування, наприклад 4-бітове, іноді може привести до погіршення точності порівняно з 8-бітовим або bfloat16. Існує компроміс між економією пам’яті та точністю, який користувачі повинні оцінити для свого випадку використання.

Квантування – це потужна техніка, яка може дозволити розгортання LLM у середовищах з обмеженими ресурсами, таких як хмарні інстанси, пристрої краю або навіть мобільні телефони, суттєво зменшуючи відбиток пам’яті.

Оцінка пам’яті для налаштування

Хоча квантування в основному використовується для ефективного висновку, техніки, такі як паралелізм тензорів і паралелізм моделей, є важливими для управління вимогами до пам’яті під час навчання або налаштування великих мовних моделей.

Пікове споживання пам’яті під час налаштування зазвичай у 3-4 рази вище, ніж під час висновку, через додаткові вимоги до пам’яті для:

• Градієнтів
• Станів оптимізатора
• Активацій з прямого проходу, збережених для зворотного проходу

Консервативна оцінка полягає в тому, що налаштування LLM з X мільярдами параметрів вимагає приблизно 4 * (2X) = 8X ГБ відеопам’яті у форматі bfloat16.

Наприклад, налаштування моделі LLaMA з 7 мільярдами параметрів вимагатиме приблизно 7 * 8 = 56 ГБ відеопам’яті на графічний процесор у форматі bfloat16. Це перевищує пам’ять сучасних графічних процесорів, що вимагає розподілених технік налаштування.

Розподільчі техніки налаштування

Було запропоновано кілька розподілених методів налаштування для подолання обмежень пам’яті графічного процесора для великих моделей:

1. Паралелізм даних: Класичний паралелізм даних реплікує всю модель на декілька графічних процесорів, розділяючи та розподіляючи пакунки навчальних даних. Це зменшує час навчання лінійно з кількістю графічних процесорів, але не зменшує пікову вимогу до пам’яті на кожному графічному процесорі.
2. ZeRO Stage 3: Розширена форма паралелізму даних, яка розділяє параметри моделі, градієнти та стани оптимізатора на графічні процесори. Це зменшує пам’ять порівняно з класичним паралелізмом даних, зберігаючи лише необхідні розділені дані на кожному графічному процесорі під час різних фаз навчання.
3. Паралелізм тензорів: Замість реплікації моделі паралелізм тензорів розділяє параметри моделі на рядки або стовпці та розподіляє їх на графічні процесори. Кожен графічний процесор обробляє розділену частину параметрів, градієнтів та станів оптимізатора, що призводить до суттєвої економії пам’яті.
4. Паралелізм конвеєра: Ця техніка розділяє шари моделі на різні графічні процесори/робочі, причому кожен пристрій виконує підмножину шарів. Активації передаються між робочими, зменшуючи пікову пам’ять, але збільшуючи накладні витрати на зв’язок.

Оцінка використання пам’яті для цих розподілених методів є нетривіальною, оскільки розподіл параметрів, градієнтів, активацій та станів оптимізатора варіюється між техніками. Крім того, різні компоненти, такі як тело трансформера та голова мовної моделі, можуть демонструвати різні поведінки розподілу пам’яті.

Рішення LLMem

Дослідники недавно запропонували LLMem, рішення, яке точно оцінює споживання відеопам’яті при застосуванні розподілених методів налаштування до LLM на декілька графічних процесорів.

Оцінка використання відеопам’яті для налаштування попередньо натренованих LLM

LLMem враховує фактори, такі як рекомбінація параметрів перед обчисленням (ZeRO Stage 3), збирання виводу у зворотному проході (паралелізм тензорів) та різні стратегії розподілу пам’яті для тіла трансформера та голови мовної моделі.

Експериментальні результати показують, що LLMem може оцінювати пікове використання відеопам’яті для налаштування LLM на одному графічному процесорі з помилкою до 1,6%, перевершуючи середню помилку DNNMem на рівні 42,6%. При застосуванні розподілених методів налаштування до LLM з понад мільярдом параметрів на декілька графічних процесорів LLMem досягає середньої помилки 3,0%.