Near Data Processing

14 February 2023

Вычисления “рядом с данными”

А теперь к заметке, с которой я сперва хотел начать этот блог. Одна из главных проблем в computer science — гигантская разница по скорости работы процессора и памяти. Это заставляет прибегать к огромному количеству ухищрений для ускорения вычислений: параллелизация, конвейеризация, спекулятивное выполнение, а также иерархия кэшей и устройств памяти. Благодаря огромной работе в области компиляторов и процессоров огромный пласт задач, где узким местом является скорость процессора, оптимизируется очень качественно, пусть и не всегда автоматически. Несколько хуже дело обстоит с вычислениями, упирающимися в скорость доступа к памяти.

Один из примеров таких вычислений — задачи машинного обучения, в частности, связанные с рекомендательными моделями. Одна из самых затратных по ресурсам частей этой задачи — так называемый инференс. Основная задача этой части максимально простая: есть массив данных фиксированной длины и есть набор индексов. Нужно извлечь из массива все данные по индексам и сложить результаты. Например, если в массиве лежат значения [0, 7, 4, 5, 12, 9], то результатом операции SLS (она же sparse length sum) с индексами (1, 3, 4) будет 7 + 5 + 12 = 24.

У этой операции есть одно очень неприятное свойство: как только набор индексов становится достаточно “разреженным”, её скорость определяется в большей степени скоростью доступа к данным. Иначе говоря, на больших таблицах с разреженными индексами (а это типичная нагрузка для такого типа моделей) даже с суперскалярным процессором и оптимизированным компилятором кодом инференс всё равно будет тормозить из-за медленного доступа к памяти (на порядки медленнее скорости работы самого процессора).

Что можно делать с такими задачами? Один из вариантов предложили Facebook в своём RecNMP или Samsung со своим AxDIMM (статья за пейволлом, но в открытом доступе есть лекция на тему). Идея их подхода состоит в том, что на плашке памяти рядом с самими ячейками памяти ставится относительно простой вычислитель, который умеет считать SLS на каком-то своём подынтервале. Затем эти значения агрегируются уже в конечный результат.

Этот подход позволяет значительно сэкономить на перемещении данных между процессором и оперативной памятью: большая часть вычислений происходит непосредственно в RAM. Улучшения при таком подходе вполне подтверждаются и практикой: тот же AxDIMM по сравнению с обычной памятью позволяет увеличить пропускную способность в полтора раза! Короче говоря, тема очень интересная и перспективная, особенно если учесть, что применять это можно далеко не только в машинном обучении.