Ускорение генерации токена LLM в два раза для больших контекстов
Помимо ChatGPT и многочисленных конкурентов в облаке с веб-мордами и/или API, существует огромная экосистема для запуска LLM на собственном железе. На Huggingface на любой бюджет найдется модель для скачивания, которая влезет в видеопамять (или в RAM, можно и на CPU запускать, если пользователь терпеливый). Вчера здесь на Хабре была очень неплохая обзорная статья.
Самые популярные open source тулы для локального запуска LLM — llama.cpp и vllm (и их многочисленные обертки). У них немного разные ниши, и дальше я буду писать о llama.cpp. Она поддерживает все возможные комбинации железа и ОС — Linux, MacOS, Windows; x86 CPU, Arm, Apple Silicon CPU & GPU, Nvidia, AMD, … Но автор и мейнтейнер — Георгий Герганов использует для разработки Mac Studio. Почему такой выбор железа?
Производительность генерации каждого токена LLM в одном потоке ограничена вычислительной мощностью в процессе построения KV-кэша (анализ промпта до генерации первого токена), и пропускной способностью памяти при генерации последующих токенов. При этом в обоих случаях очень полезно уметь быстро загружать веса из видеопамяти в ALU видеокарты (или CPU).
А на каких платформах у нас самая толстая труба от памяти к ALU? На Nvidia, потом AMD GPU, и на почетном третьем месте — Apple SoC GPU. Это просто сортировка по ширине трубы. Но чтобы запустить действительно большую модель, необходимо, чтобы она влезла в память. Apple (за негуманную цену) продает ноутбуки с 128 гигабайтами «видеопамяти», и Mac Studio с 192 гигабайтами. Но если покупать видеокарты с таким же объемом видеопамяти, то дешевле не выйдет. (Кстати, если видео памяти не хватает на одной машине/GPU, то можно добавить еще железа, масштабируется по горизонтали оно практически линейно. Так что Llama 3.1 700G можно запустить с 4b квантованием на четырех Macbook Pro максимальной конфигурации, будет генерировать несколько десятков токенов в минуту)
Рассмотрим поподробнее, на что тратится время при генерации токена на Apple silicon GPU. При использовании KV-кэша, самой дорогой операцией в каждом слое трансформера будет вычисление attention, и в нем самой долгой операцией будет не произведение матриц (как при генерации KV-кэша), а скалярное произведение матрицы на вектор.
Причем, количество элементов вектора равно длине промпта, округленной вверх до ближайшего множителя 128, а количество строк матрицы — обычно ровно 128.
Реализация соответствующего Metal kernel для Apple GPU в llama.cpp достаточно наивна —
на каждый элемент вектора запускается по одному потоку, каждый из которых, обычно, делает всего четыре умножения. В результате, при длинном контексте, потоков у нас слишком много, а ALU в среднем загружено на 7%, так как GPU занято в-основном созданием потоков и записью результатов их вычислений в память.
Когда промпт перерастает несколько десятков тысяч токенов (а многие современные модели обучены поддерживать контекст до 128k токенов и выше), с производительностью становится все очень плохо. Я обнаружил это в феврале, открыл баг в марте, и Георгий ответил, что было бы неплохо это исправить. Фикс был простой и очевидный — запускать в 32 раза меньше потоков, и давать каждому потоку в 32 раз больше работы. Но отлаживал этот Metal kernel код я слишком долго, и предложил патч только в конце мая (некоторое время еще ушло на то, чтобы мой тогдашний работодатель официально разрешил послать этот патч).
Ускорение получилось в 2 и более раз для некоторых моделей и очень больших длин контекста. Но в результате, патч не приняли, так как незадолго до этого в llama.cpp Metal был реализован механизм flash attention, который ускоряет вычисление всего блока attention трансформера, а не только лишь произведения матрицы на вектор. Эта мега-оптимизация представляет собой один супер-длинный Metal kernel, который включает в себя все операции attention. Работает, как правило, немного быстрее, чем старый код с моим патчем (правда, по умолчанию в llama.cpp flash attention пока еще, вроде, не включен).
У реализации Flash Attention в llama.cpp есть другая проблема — он запускает только 32 групп потоков, и следовательно, на Apple Silicon GPU использует только 32 ядра. Так что на машинах с большим количеством GPU ядер некоторая их часть вообще не используется. Но сейчас у меня уже нет доступа к неограниченному количество high end железа Apple, поэтому поправить этот performance bug я уже не могу. Так что публикую этот кейс на Хабре, и желаю удачи, если найдутся желающие еще ускорить inference llama.cpp на Apple GPU.
