Увеличение скорости генерации токенов LLM на Apple Silicon

Введение

Существует широкий спектр возможностей для запуска больших языковых моделей (LLM) на собственном оборудовании, помимо привычных вариантов использования облачных сервисов и API. На платформе Huggingface доступны модели различных бюджетов, позволяя пользователям найти оптимальное решение для своих потребностей. Наличие моделей, поддерживающих работу с видеопамятью, RAM или даже CPU, расширяет возможности выбора и адаптации их к индивидуальным требованиям и особенностям задачи. Важно отметить, что с появлением новых технологий и оптимизаций, процесс запуска моделей на собственных устройствах становится все более эффективным и доступным для широкого круга пользователей.

Популярные инструменты

Георгий Герганов, автор и мейнтейнер llama.cpp, предпочитает использовать Mac Studio для разработки своего продукта. Он делает это из-за важных технических причин, связанных с увеличением скорости генерации токенов LLM на платформе Apple Silicon. На Apple SoC GPU имеется широкая труба от памяти к ALU, что существенно способствует быстрой загрузке весов из видеопамяти в ALU видеокарты. Для работы с моделями, требующими большой памяти, Apple предлагает устройства с впечатляющими объемами видеопамяти - ноутбуки с 128 гигабайтами или Mac Studio с 192 гигабайтами, что делает их привлекательным выбором при ограничениях бюджета. Георгий даже укрепил производительность llama.cpp на Apple Silicon GPU, исправив оптимизационные ошибки в программном коде, что привело к значительному ускорению работы для определенных моделей. В итоге, выбор Mac Studio для разработки llama.cpp отчасти определяется исключительными техническими возможностями и особенностями платформы Apple Silicon.

Аппаратные ограничения

Производительность генерации токенов LLM на Apple Silicon GPU ограничена вычислительной мощностью при построении KV-кэша и использовании памяти для генерации последующих токенов. Загрузка весов из видеопамяти в ALU видеокарты играет ключевую роль. Наилучшими платформами для этого являются Nvidia, AMD GPU и Apple SoC GPU. Однако, для работ с крупными моделями, важно, чтобы они помещались в память. Apple предлагает устройства с высоким объемом памяти, включая ноутбуки с 128 гигабайтами «видеопамяти» и Mac Studio с 192 гигабайтами. Приобретение видеокарт с таким же объемом памяти представляется менее экономичным решением.

Время генерации токена на Apple silicon GPU

При генерации токена на Apple Silicon GPU основное время уходит на операции, связанные с вычислением attention в каждом слое трансформера. Наиболее ресурсоемкой операцией является скалярное произведение матрицы на вектор. В этом случае количество элементов вектора соответствует длине промпта, округленной до ближайшего множителя 128, а количество строк матрицы обычно равно 128. Реализация соответствующего Metal kernel для Apple GPU в llama.cpp показала, что изначально было использовано слишком много потоков, что не позволяло эффективно загружать вычислительные ресурсы. Однако после внесения патча, ускорение работы моделей на GPU составило два и более раза, что является значительным улучшением. Однако в последствии был разработан механизм flash attention, который дополнительно оптимизировал вычисления блока attention, ускоряя процесс генерации токенов.

Реализация и оптимизация

При создании Metal kernel для Apple GPU в llama.cpp выявлены недочеты в эффективности потребления ресурсов. В частности, каждый элемент вектора запускается на свой поток, что приводит к избыточному количеству потоков и недостаточной загрузке ALU. При обработке длинных контекстов производительность снижается, так как большая часть работы заключается в создании и записи данных в память GPU. Оптимизация кода сокращением числа потоков в 32 раза и увеличением рабочей нагрузки каждого потока привела к значительному ускорению работы для некоторых моделей и длинных контекстов. Однако, внедрение Flash Attention в llama.cpp дало дополнительный прирост скорости вычислений за счет оптимизированного механизма attention, хотя существует проблема использования не всех доступных ресурсов ядер GPU на машинах с большим количеством ядер.

Предложенные патчи

Обнаружив проблему в работе генерации токенов на Apple Silicon GPU, я предложил патч для исправления этой проблемы. Мои улучшения заключались в запуске меньшего количества потоков, каждому из которых предоставлялось больше работы, что привело к ускорению процесса в несколько раз. Однако, хотя мой патч дал положительные результаты, его не приняли из-за внедрения в llama.cpp механизма flash attention, который предоставлял еще более значительное улучшение в скорости работы блока attention трансформера. Желаю успехов тем, кто продолжит работу над улучшением производительности llama.cpp на Apple GPU.

Flash Attention

Оптимизация Flash Attention в llama.cpp представляет собой единый Metal kernel, интегрирующий все операции attention. По умолчанию этот механизм работает немного быстрее предыдущих версий кода, включая патч автора. Однако, существует ограничение - при выполнении операции Flash Attention используются только 32 группы потоков, что означает активацию лишь 32 ядер на графическом процессоре Apple Silicon. На машинах с обильным количеством ядер GPU, остальные не задействованы, оставаясь неиспользованными. Этот фактор можно рассматривать как некоторое ограничение использования потенциала аппаратного обеспечения, требуя дальнейших улучшений в эффективности распределения вычислительных ресурсов для повышения производительности и оптимизации работы LLM на устройствах с Apple Silicon GPU.

Заключение

Похоже, что процесс развития и совершенствования программного обеспечения, специализированного на ускорении работы с нейросетями, постоянно представляет собой балансировку между ресурсами и оптимизацией. История исправления performance bug в llama.cpp на Apple GPU наглядно показывает, как даже кажущиеся незначительные детали в коде могут влиять на общую производительность и эффективность работы модели. Важно помнить о постоянной необходимости апгрейда и оптимизации кода для улучшения результатов и избежания возможных узких мест в процессе генерации токенов. Мир разработки искусственного интеллекта и машинного обучения постоянно меняется, требуя от специалистов готовности искать новые пути для оптимизации и улучшения производительности.