Pytorch introduces TorchRec, an open source library to build recommendation systems

CEO Mark Zuckerberg has introduced TorchRec, an open source library for building state-of-the-art recommendation Systems under PyTorch, at Inside the Lab event. The new library provides common sparsity and parallelism primitives, enabling researchers to build state-of-the-art personalisation models and deploy them in production. It  includes a scalable low-level modelling foundation alongside rich batteries-included modules.

Meta initially targeted“two-tower” architectures with separate submodules to learn representations of candidate items and the query or context. Input signals can be a mix of floating point “dense” features or high-cardinality categorical “sparse” features that require large embedding tables to be trained. Efficient training of such architectures involves combining data parallelism that replicates the “dense” part of computation and model parallelism that partitions large embedding tables across many nodes.

The library includes optimised Recommendation Systems kernels that run on FBGEMM, a high-performance kernel library, modelling primitives such as jagged tensors and embedding bags, to create multinodal models using model parallelism. The PyTorch team has released TorchRec after close to two years in testing.

TorchRec was used to train a model with 1.25 million parameters that went into production in January and a 3 trillion parameter model that is expected to go into production.