Управление распределением блоков по серверам для децентрализованного инференса LLM модели в системе Petals

В современном мире огромной популярностью пользуются большие языковые модели (LLM), которые часто имеют более 100 миллиардов параметров. Сегодня все желающие могут бесплатно загрузить предварительно обученные модели такого масштаба, однако для использования подобных моделей требуется высокотехнологичное и дорогое оборудование, недоступное многим исследователям. Одним из способов решения этой проблемы является PETALS - система для коллаборативной работы с большими NLP моделями, которая объединяет ресурсы нескольких участников в один пул. Одним из основных сценариев использования системы является инференс (вывод) обученной модели для получения предсказаний пользователями. В текущей реализации PETALS стратегия аллокации блоков, которые будет обслуживать сервер, оптимизирует вычислительную пропускную способность на каждом слое модели. Этот подход хорошо показывает себя для обучения, но не для инференса. В этой работе было проведено исследование задачи распределения блоков, оптимизирующий время клиентских запросов на инференс обученной модели. С помощью рассмотренных алгоритмов получилось ускорить вывод большой языковой модели до 1.8 раз, по сравнению с базовой стратегией в тестовом симуляторе.