Нейролента - подборка новостей о нейронных сетях, ChatGPT

PowerInfer: Fast Large Language Model Serving with a...

PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU
Yixin Song, Zeyu Mi, Haotong Xie, Haibo Chen
Статья: https://arxiv.org/abs/2312.12456
Код: https://github.com/SJTU-IPADS/PowerInfer

Очень интересный инженерный проект по оптимизации инференса LLM от Института параллельных и распределённых систем из Шанхая, чтобы запускать большие LLMки на GPU консьюмерского уровня. Авторы опубликовали свой гибридный GPU-CPU движок для инференса под названием PowerInfer, который работает быстрее llama.cpp в 11+ раз на RTX 4090 (24G) с моделью Falcon(ReLU)-40B-FP16, а также демонстрирует скорость генерации токенов всего на 18% ниже, чем в 10 раз более дорогая A100.

Работа основана на следующих наблюдениях. Первое: активации нейронов в LLM демонстрируют локальность (locality) и распределены по степенному закону. Есть небольшое количество hot neurons, которые активируются часто на разных входных данных, и cold neurons, активирующиеся редко на конкретных входах, и которых большинство. Второе ценное наблюдение в том, что пересылать нужные веса с CPU на GPU и считать результат там может быть медленнее, чем посчитать их на месте на CPU, особенно если нужна не полная матрица весов, а её подмножество (которое может быть неплохо посчитано через современные векторные расширения в системе команд процессоров). В итоге горячие нейроны (строки/столбцы соответствующих матриц весов) отправляются на GPU для быстрого доступа, а холодные считаются на CPU по мере надобности. Это существенно понижает требования к памяти GPU и уменьшает пересылку данных через шину PCIe туда-сюда.

Технически PowerInfer состоит из двух частей: движка инференса и оффлайн части.

Оффлайн часть включает профайлер для обнаружения горячих нейронов и солвер (на основе целочисленного линейного программирования, Integer Linear Programming, ILP) для оценки влияния нейрона на результат инференса и создания полиси аллокации нейронов. Это 400 строк на Питоне + HuggingFace Transformers. Горячие/холодные нейроны определяются на офлайн шаге, далее во время инференса они будут загружены на GPU/CPU соответственно.

Движок инференса сделан на базе llama.cpp с добавленными 4200 строк кода на C++/CUDA. Движок запускает на CPU отдельные треды для GPU и CPU экзекьюторов, которые управляют этими конкурентными вычислениями, а слияние результата происходит на GPU (потому что там сидят более часто активирующиеся нейроны). PowerInfer использует адаптивные предикторы для предсказания, какие нейроны будут активны при текущих входных данных, и только эти нейроны (из уже разделённых на горячие и холодные) и вычисляются. То есть нейрон мог быть определён как горячий на офлайн шаге, загружен на GPU, но всё равно не активирован, потому что предиктор решил, что его вычислять не нужно. Можно, наверное, это рассматривать как adaptive sparsity.

Для каждого трансформерного слоя работают два предиктора, один для MHSA, другой для MLP. Живут они тоже в памяти GPU, поэтому есть конкуренция за этот ресурс и предикторы важно сделать минимально допустимого размера. Предиктор -- это MLP с одним скрытым слоем. В работе есть хитрая итеративная процедура для изменения размера скрытого слоя предиктора в зависимости от асимметрии (skewness) распределения активаций.

Так понимаю, что предикторы и для hot/cold и для онлайн шага обучены независимо на типовых корпусах типа C4 или Википедии, и код их обучения пока не опубликован. Но вроде как планируется.

Также реализованы свои neuron-aware sparse operators, которые более заточены на эту задачу, чем более универсальные из cuSPARSE, Sputnik (https://github.com/google-research/sputnik), SparTA (https://github.com/microsoft/SparTA) или FlashLLM (https://github.com/AlibabaResearch/flash-llm). А JIT компилятор PIT (https://arxiv.org/abs/2301.10936) не поддерживает GPU-CPU гибридизацию.

В результате хорошие ускорения до десятка раз относительно llama.cpp на разных моделях.