Большие языковые модели с успехом применяются в биоинформатике. Неудивительно, ведь язык биологических последовательностей – всего лишь часть множества знаковых систем, существующих в природе. Но есть нюанс: осмысленные фразы языка биоинформатики, как правило, гораздо длиннее фраз языков межчеловеческого общения. Популярные языковые модели на основе трансформеров требуют квадратичного увеличения объема вычислений с ростом длины входной последовательности.

Мое внимание привлекла модель Mamba, демонстрирующая всего лишь линейный рост сложности с увеличением длины ввода, и из-за этой особенности прекрасно подходящая для задач биоинформатики. За деталями отправляю вас к оригинальной статье на Arxiv и двум публикациям на Хабре, здесь и здесь.

Официальная версия Мамбы оптимизирована под GPU с использованием технологии CUDA. К сожалению, все графические ускорители в облаке Microsoft Azure уходят коммерческим пользователям. А исследователям Microsoft Research оборудование достается по остаточному принципу. Я решил оставить ценный ресурс своим коллегам и запустить свою модель на CPU, предварительно оптимизировав ее под векторные вычисления.

Чтобы не заниматься отколупыванием CUDA кода от модели, авторы оригинальной статьи посоветовали мне начать с чисто Питонового кода из вот этого репозитория.

Прежде чем заняться оптимизацией, соберем профиль выполнения модели. Ограничим наш анализ первыми пятью строками отчета; остальное, по сути, шум, мало сказывающийся на производительности.

Строка 2 — линейная проекция в начале и в конце Mamba блока, по сути, матричное умножение. Строка 6 — одномерная свертка. Мы оставим их в покое и используем в качестве эталона времени. А строки с 3 по 5 — это ядро алгоритма Mamba, в сумме эти строки отнимают более половины времени выполнения кода. Сосредоточимся на оптимизации именно этих строк, составляющих метод selective_scan.

Все самое интересное происходит в вызове функции einsum, то есть, обобщенного тензорного перемножения в нотации Эйнштейна. Разумеется, функции библиотеки einops, как и весь PyTorch, написаны на C и уже оптимизированы талантливыми программистами. Это весьма распространенная ситуация, когда я не начинаю работу по оптимизации с чистого листа, а имею для сравнения результат работы опытного разработчика, использующего оптимизирующий компилятор. Посмотрим, удастся ли нам превзойти их результат.

Начнем нашу работу по оптимизации, как ни странно, с ухудшения производительности нашего кода. Вспомнив курс теоретической физики, распишем эйнштейновское перемножение тензоров из строки 5 в виде вложенных циклов на чистом Питоне.

Поступим аналогично со строкой 7.

Строки с 9 по 17 также могут быть расписаны в виде вложенных циклов.

Полученный код выполняется чудовищно медленно, ведь мы отказались от оптимизированных библиотек в пользу чистого Питона. Зато теперь лучше видна структура кода, позволяющая инвариантные преобразования.

Во-первых, заметим, что все три вышеприведенных четверных цикла имеют одни и те же верхние границы и могут быть слиты в один. При этом оказывается, что промежуточные тензоры deltaA и deltaB_u не нужно сохранять в памяти, а можно использовать вычисленные значения немедленно.

Во-вторых, изменим порядок циклов с h, i, j, k на h, j, k, i. Это допустимое инвариантное преобразование программы.

>

Получилась очень интересная конструкция. Во внутреннем цикле вычисляются элементы xk, каждый из которых зависит от предыдущего, и последовательно прибавляются с некоторым весом к накопительной сумме; по-английски это операция называется scan. Именно эту операцию накопительной суммы с весом авторы метода и назвали selective scan. Несмотря на взаимозависимость членов накопительной суммы, ее возможно, пусть и с некоторыми ухищрениями, вычислять параллельно, и векторизация накопительной суммы – мой конек. Тем более, что учебник предписывает оптимизировать, в нашем случае векторизовать, самый внутренний цикл.

Но мы опять поступим неправильно и оставим внутренний цикл в покое.

Второй изнутри цикл мы тоже не станем векторизовывать. Поскольку этот цикл короткий, параметр модели channel_count всего 16, то этот цикл мы размотаем в последовательный код, подбросив процессору возможностей для внеочередного выполнения операций.

А вот третий изнутри цикл обладает чрезвычайной параллельностью, то есть вычисления на каждой итерации цикла не зависят от прочих итераций. Идеально подходит для векторизации.

Несколько замечаний. Во-первых, здесь нам понадобилась функция векторного вычисления экспоненты, подробно разобранная в отдельной публикации. Тема интересная, но слабо пересекающаяся с оптимизацией языковой модели.

Во-вторых, тензор А во входных данных транспонирован по сравнению с исходной функцией. Мне было удобнее транспонировать данные в TensorFlow, чем усложнять векторный код.

И, наконец, в строке 33 я вспомнил, что, кроме оптимизации под одно ядро, можно использовать возможности многоядерной системы. Поскольку мы выбрали для оптимизации чрезвычайно параллельный цикл, разложить его по ядрам не составило труда. В этом мое кредо: прежде чем задействовать многие ядра и многочисленные ресурсы в облаке, следует качественно оптимизировать программу для выполнения на одном ядре.

Измерим производительность оптимизированной версии.

Строка 4 нового отчета заменила собой строки 3-5 оригинала. Время выполнения функции уменьшилось почти в 20 раз. Неплохой результат.

Следующий логичный шаг — попытка ускорить перемножение матриц в операторе линейной проекции. Но это уже совсем другая история.