meta-pytorch/torchrec

What it solves

TorchRec addresses the challenges of training and deploying large-scale recommender systems (RecSys). These models typically use high-cardinality categorical sparse features that require massive embedding tables, which are often too large to fit on a single GPU, necessitating efficient distribution across multiple devices and nodes.

How it works

It provides a set of sparsity and parallelism primitives built on PyTorch. The library includes:

• Sharding Strategies: Tools to split embedding tables using various methods (e.g., table-wise, row-wise, column-wise) to distribute the model across GPUs.
• Automatic Planning: A planner that generates optimized sharding plans to maximize performance.
• Pipelined Training: A system that overlaps data loading, device transfers, and computation to reduce bottlenecks.
• Optimized Kernels: High-performance RecSys kernels powered by FBGEMM.
• Quantization: Support for reduced precision to speed up training and inference and optimization for C++ inference.

Who it’s for

It is designed for researchers and engineers building state-of-the-art personalization and recommendation models that need to scale to many GPUs and production environments.

Highlights

• Powers production RecSys models at Meta.
• Supports hybrid data-parallelism and model-parallelism.
• Includes built-in RecSys datasets like Criteo click logs and MovieLens.
• Provides end-to-end training examples, such as the DLRM event prediction model.