Как запустить большую языковую модель OPT-175B, аналог GPT-3 на своём компьютере? 🔥 FlexGen !

Недавно состоялся релиз библиотеки на PyTorch FlexGen, которая позволяет запустить огромную языковую модель типа OPT-175B/GPT-3 на одной видеокарте (16GB T4 GPU или 24GB RTX3090). Правда нужно много RAM (200Gb++).

FlexGen, это еще один инструмент для запуска больших языковых моделей на одной видеокарте. Работает за счет сжатия весов и offloading'а, наподобие HuggingFace Accelerate и DeepSpeed Zero, но, производительнее на порядок.

Так открытую языковую модель OPT-30, теперь можно запустить на одной видеокарте (16GB GPU) со скоростью ~1 токен в секунду!

Для ускорения инференса модели достаточно написать:
model = OptLM(model)

Для достижения такой оптимизации используется несколько приёмов, такие, как гибкое автоматическое распределение весов на GPU, CPU или SSD диск, сжатие весов до 4 бит вместо стандартных для языковых моделей 16 бит с минимальной потерей точности, и блочный проход вместо принятого для других оптимизаторов построчного, что позволяет эффективнее использовать пропускную способность памяти и реже перегружать тензоры между различными уровнями кеша.

В итоге производительность увеличилась до 100 раз по сравнению с другими системами оптимизации и ускорения при использовании больших языковых моделей.

GitHub: https://github.com/FMInference/FlexGen

Whitepaper: https://github.com/FMInference/FlexGen/blob/main/docs/paper.pdf

В качестве теста я запустил OPT-1.3B на своём компьютере:

(venv) D:\Projects\FlexGen>python -m flexgen.flex_opt --model facebook/opt-1.3b NVIDIA GeForce RTX 3080 model size: 2.443 GB, cache size: 0.398 GB, hidden size (prefill): 0.008 GB warmup - init weights warmup - generate benchmark - generate benchmark - delete weights Outputs: ---------------------------------------------------------------------- 0: Paris is the capital city of France. It is the most populous city in France, with an estimated population of 6,848,000 in 2016. It is the second most populous city ----------------------------------------------------------------------

TorchDevice: cuda:0 cur_mem: 0.0000 GB, peak_mem: 3.2399 GB TorchDevice: cpu cur_mem: 0.0000 GB, peak_mem: 0.0000 GB model size: 2.443 GB cache size: 0.398 GB hidden size (p): 0.008 GB peak gpu mem: 3.240 GB projected: False prefill latency: 0.318 s prefill throughput: 6435.642 token/s decode latency: 1.134 s decode throughput: 109.319 token/s total latency: 1.453 s total throughput: 88.122 token/s

Как видите, модель в 1.3 миллиарда токенов отработала за полторы секунды!