Сегодня разберём работу Alibaba под названием RecIS (Recommendation Intelligence System). Это фреймворк для обучения рекомендательных систем, который нужен, чтобы масштабировать «всё и вся» в разных контекстах, оставаясь при этом PyTorch-френдли и пригодными для продакшена.
В рексистемах есть два разных типа нагрузки. В sparse-части, связанной с эмбеддингами, мало вычислений, но много обращений к памяти. В dense-части, наоборот, много вычислений и меньше работы с памятью. В рексистемах по сравнению с другими областями (NLP, CV) ощутимо больше sparse-нагрузки, поэтому оптимизации касаются именно этой части.
В статье обращают внимание, что скорость доступа к памяти на GPU быстро растёт и уже сильно выше, чем на CPU. Долго считалось, что доступ к памяти остаётся боттлнеком и на GPU, но в работе говорят, что память уже не такая и медленная. Отсюда идея, что можно сделать размен в другую сторону — больше обращений к памяти и уменьшение количества вычислений за счёт большего количества операций с ней.
Авторы вводят метрику MBU — максимальная утилизация пропускной способности. Опираются на идею roofline-модели, где оценивается баланс между вычислениями и доступом к памяти. В классическом варианте по горизонтали откладывают число операций на загруженный байт, а по вертикали — количество FLOPs, которые модель может утилизровать.
Для рексистем такая постановка не очень подходит из-за sparse-нагрузки. Важно не количество вычислений, а то, насколько хорошо утилизируется пропускная способность памяти. Поэтому в работе рассматривают bandwidth-based-вариант, где по вертикали откладывается bandwidth, которую может утилизировать модель, а по горизонтали — объём обращений к памяти на количество вычислений.
Одно из применений этой идеи — работа со sparse-эмбеддингами. Стандартно таблица эмбеддингов реплицируется на все GPU: после каждой эпохи градиенты собираются, усредняются и рассылаются обратно. Здесь предлагают шардировать таблицу так, чтобы каждая GPU хранила только свою часть.
На форварде каждая GPU определяет, какие эмбеддинги ей нужны с других устройств, после происходит all-to-all-коммуникация, и карты обмениваются нужной информацией. Затем каждая GPU читает свои эмбеддинги и отправляет их туда, где они нужны.
На бэкварде происходит то же самое: каждая GPU знает, куда нужно записать градиенты, и они пересылаются напрямую в соответствующие шарды. В итоге всё укладывается в две all-to-all-коммуникации, а таблица занимает меньше памяти на каждой видеокарте.
Также делают оптимизации. Есть стандартные вещи, такие как mixed precision, FlashAttention, ZeRO, fused softmax и cross-entropy. Есть специфичные для рекомендаций. Например, если у нескольких embedding-таблиц одинаковая размерность, их объединяют в одну, чтобы оптимизировать доступ к памяти.
Данные хранят в колоночном формате, для sparse-структур используют CSR.
После загрузки данных есть всего четыре типа операций:
- бакетизация / взятие по модулю;
- (мульти)хэширование;
- обрезание истории;
- построение пар фичей.
Каждую фичу представляют как последовательность таких операций. Некоторые из них (например, хэширование по нескольким разным сидам) можно объединить и зафьюзить.
Эксперименты проводят на двух моделях: MSE (предсказание вероятности клика) и LMA (предсказание следующего баннера по истории). По времени обучения RecIS работает быстрее, чем TensorFlow и PyTorch. По метрике MBU тоже более высокая утилизация пропускной способности.
При этом упоминают и проблемы:
- загрузка данных из сети может занимать много времени (сэмплы до ~6GB);
- embedding-таблицы остаются очень большими, в том числе из-за старых объектов;
- в старых фреймворках вроде TensorFlow 1.x есть ограничения на размер тензоров.
В конце авторы делают вывод, что в современных multi-GPU-системах памяти обычно достаточно. И это позволяет сместить баланс: делать больше операций с памятью и меньше вычислений.
@RecSysChannel
Разбор подготовил
