Анонс библиотеки TorchRec для создания рекомендаций

Команда PyTorch анонсировала выпуск библиотеки TorchRec для создания рекомендательной системы. Она включает примитивы разреженности и параллелизма, которые позволяют строить новейшие рекомендательные модели и разворачивать их в production.

С чего начиналось

Многие продуктовые компании пользуются рекомендательными системами (Recommendation Systems, RecSys). Но каждый знает о них в виде проектов с открытым исходным кодом. В отличие от сфер Computer Vision и NLP, в большинстве случаев разработка рекомендательных систем ведется за закрытыми дверями компании. Причем то же самое обнаруживается и с научными публикациями — мало, кто пишет о RecSys. При том что, RecSys основана на обучении моделей с использованием разреженных и/или последовательных событий, что имеет общие пересечения с другими сферами Machine Learning. Многие из этих методов переносимы, в частности, для масштабирования и распределенного выполнения. Большая часть глобальных инвестиций в ИИ направлена на разработку этих методов RecSys, поэтому их сокрытие от глаз сообщества блокирует инвестиции на развитие более широких сфер ИИ.

В середине 2020 года команде PyTorch пришло много сообщений о том, что неплохо было бы иметь пакет для построения рекомендательной системы с открытым исходным кодом в составе экосистемы PyTorch. В это же время время инженеры из Meta (ex-Facebook) изъявили свое желание помочь развитию рекомендательных систем в виде библиотеки, которая будет постоянно поддерживаться и развиваться. Это, действительно, была хорошая идея, от которой получат пользу как научные сотрудники, так и компании. Поэтому команда PyTorch начала наращивать кодовую базу такой системы. После двух лет написания кода, оценок производительности, миграций и тестирования в мире Open Source появился TorchRec.

Что есть в TorchRec

TorchRec включает масштабируемую, низкоуровневую основу вместе со множеством дополнительных модулей. Изначально планировалось построить двубашенную (two-tower) архитектуру ([1], [2]), которая имеет раздельные модули для обучения представлений и запросов или контекстов. Входные сигналы могут состоять из плотных или разреженных массивов, что требует огромные встроенные таблицы для обучения. Эффективное обучение такой архитектуры требует объединения параллелизма данных, что воспроизводит “плотную” часть вычислений и параллелизма, который использует партиции больших встроенных таблиц по многим узлам.

В частности, TorchRec включает:

• Примитивы моделирования, например, embedding bags и jagged tensors, которые позволяют разрабатывать большие, производительные, мультиаппаратные/мультиузловые модули с использованием гибридного параллелизма данных и параллелизма модели.
• Оптимизированные ядра RecSys на основе FBGEMM, включая поддержку разреженности и квантованных операций.
• Возможность создавать партиции встроенных таблиц с большим количеством вариаций: табличным, построчным, табличнострочным, колоночным сегментированием.
• Создание конвейеров с использованием перемещений данных (копии в GPU), межаппаратных коммуникаций (input_dist) и вычислений (прямого и обратного распространения) для повышения производительности.
• Поддержка получения выводов GPU.
• Общие модули для RecSys, например, модели и датасеты (Criteo & Movielens).

Чтобы показать гибкость RecSys, ниже приведен код на Python модели DLRM.

# Задание разреженных встроенных слоев
eb_configs = [
   EmbeddingBagConfig(
       name=f"t_{feature_name}",
       embedding_dim=64,
       num_embeddings=100_000,
       feature_names=[feature_name],
   )
   for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
]

# Импорт и инстанциирование модели со "встраиваемой" конфигурацией
# Устройство `meta` обозначает ленивое инстанциирование, без выделения памяти
train_model = DLRM(
   embedding_bag_collection=EmbeddingBagCollection(
       tables=eb_configs, device=torch.device("meta")
   ),
   dense_in_features=len(DEFAULT_INT_NAMES),
   dense_arch_layer_sizes=[512, 256, 64],
   over_arch_layer_sizes=[512, 512, 256, 1],
   dense_device=device,
)

# Распространение модели по девайсам, так же как и в DDP
model = DistributedModelParallel(
   module=train_model,
   device=device,
)

optimizer = torch.optim.SGD(params, lr=args.learning_rate)
# Оптимизация модели в цикле, так же как это делается с другими моделями
# Или можно использоват конвейер для синхронизации коммуникаций и вычислений
for epoch in range(epochs):
   # Обучение

Масштабируемая производительность

TorchRec имеет новейшую (state-of-the-art) инфраструктуру для создания масштабируемой рекомендательной системы с самыми крупными моделями Meta. Ее использовали для обучения модели с 1.25 триллионов параметрами, которая была выпущена в production в январе. Модель с 3 триллионами параметрами скоро также будет доступна. Это показывает, что PyTorch полностью способен решать проблему масштабирования рекомендательной системы. Правда, есть одно “но” — очень трудно провести оценки производительности крупной модели на публичном датасете из-за того, что многие инструменты оценки с открытым исходным кодом слишком малы.

О том, как строить свои рекомендательные системы вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и ИТ-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

• PYML: Машинное обучение на Python
• Разработка и внедрение ML-решений
• Нейронные сети на Python
• NLP с Python
• Computer vision на Python

1. Оригинал