INTUNE: REINFORCEMENT LEARNING-BASED DATA PIPELINE OPTIMIZATION FOR DEEP RECOMMENDATION MODELS
https://arxiv.org/abs/2308.08500 #recsys23

IT-гиганты используют большие нейросети (DLRM), чтобы строить персонализированные рекомендации для сотен миллионов пользователей. Количество обучаемых параметров у таких моделей сравнимо с количеством параметров у современных языковых моделей. При этом DLRM намного проще трансформеров с точки зрения архитектуры. Поэтому при обучении DLRM узкое место — не вычисления внутри нейронной сети, а процесс подготовки данных, который состоит из нескольких этапов:
- загрузки данных с диска;
- формирования батчей;
- перемешивания;
- преобразований данных;
- загрузки данных на GPU.
Если хотя бы один из этапов «тормозит» — весь процесс замедляется, а вместе с ним замедляется и обучение. Решением в этой ситуации является распараллеливание каждого из этапов на нескольких процессорах, причем так, чтобы пропускная способность всех этапов была одинаковой.

Ключевым вопросом здесь становится выявление необходимого количества процессоров на каждом этапе. Оказывается, что для DLRM стандартные инструменты вроде AUTOTUNE работают недостаточно хорошо: процессоры используются неоптимально, переполняется память. Разработчики из Netflix предлагают собственное решение.

Представленный в статье фреймворк InTune основан на обучении с подкреплением в такой постановке:
- Action. Агент InTune может увеличивать или уменьшать количество процессоров, которое выдается каждому этапу подготовки данных.
- State. В качестве параметров среды агент наблюдает метрики процесса обучения — например, текущую пропускную способность и объём свободной памяти в кластере.
- Reward. С помощью Q-learning агент учится выделять ресурсы так, чтобы скорость обучения была максимальной при минимальном использовании памяти.

Эксперименты показали, что под управлением InTune модель обучается в 2,3 раза быстрее, чем под управлением AUTOTUNE, а доля падений по памяти снижается с 8% до нуля. Авторы не опубликовали исходный код, но идея достаточно простая, чтобы повторить её в «домашних условиях».

Бонус: предложенный в статье фреймворк может быть полезен не только для обучения DLRM. Например, его можно прикрутить для автоматического выделения ресурсов приложениям Spark.