Python-Torchrec: Pytorch domain library for recommendation systems

TorchRec (Experimental Release)

TorchRec is a PyTorch domain library built to provide common sparsity & parallelism primitives needed for large-scale recommender systems (RecSys). It allows authors to train models with large embedding tables sharded across many GPUs.

TorchRec contains:

  • Parallelism primitives that enable easy authoring of large, performant multi-device/multi-node models using hybrid data-parallelism/model-parallelism.
  • The TorchRec sharder can shard embedding tables with different sharding strategies including data-parallel, table-wise, row-wise, table-wise-row-wise, and column-wise sharding.
  • The TorchRec planner can automatically generate optimized sharding plans for models.
  • Pipelined training overlaps dataloading device transfer (copy to GPU), inter-device communications (input_dist), and computation (forward, backward) for increased performance.
  • Optimized kernels for RecSys powered by FBGEMM.
  • Quantization support for reduced precision training and inference.
  • Common modules for RecSys.
  • Production-proven model architectures for RecSys.
  • RecSys datasets (criteo click logs and movielens)
  • Examples of end-to-end training such the dlrm event prediction model trained on criteo click logs dataset.


We are currently iterating on the setup experience. For now, we provide manual instructions on how to build from source. The example below shows how to install with CUDA 11.1. This setup assumes you have conda installed.

  1. Install pytorch. See pytorch documentation
conda install pytorch cudatoolkit=11.3 -c pytorch-nightly
  1. Next, install FBGEMM_GPU from source (included in third_party folder of torchrec) by following the directions here. Installing fbgemm GPU is optional, but using FBGEMM w/ CUDA will be much faster. For CUDA 11.1 and SM80 (Ampere) architecture, the following instructions can be used:
conda install -c conda-forge scikit-build jinja2 ninja cmake
export CUB_DIR=/usr/local/cuda-11.1/include/cub
export CUDA_BIN_PATH=/usr/local/cuda-11.1/
export CUDACXX=/usr/local/cuda-11.1/bin/nvcc
python install -Dcuda_architectures="80" -DCUDNN_LIBRARY_PATH=/usr/local/cuda-11.1/lib64/ -DCUDNN_INCLUDE_PATH=/usr/local/cuda-11.1/include

The last line of the above code block (python install...) which manually installs fbgemm_gpu can be skipped if you do not need to build fbgemm_gpu with custom build-related flags. Skip to the next step if that is the case.

  1. Download and install TorchRec.
git clone --recursive

# cd to the directory where torchrec's is located. Then run one of the below:
cd torchrec
python build develop --skip_fbgemm  # If you manually installed fbgemm_gpu in the previous step.
python build develop                # Otherwise. This will run the fbgemm_gpu install step for you behind the scenes.
  1. Install torchx
pip install torchx-nightly
  1. Test the installation.
torchx run --scheduler local_cwd
  1. If you want to run a more complex example, please take a look at the torchrec DLRM example.

That's it! In the near-to-mid future, we will simplify this process considerably. Stay tuned...


TorchRec is BSD licensed, as found in the LICENSE file.


  • test (#33)
    test (#33)

    Jan 20, 2022

    Summary: Pull Request resolved:

    Differential Revision: D33677484

    CLA Signed 
  • Is there any example about how to train model with large embeddings?
    Is there any example about how to train model with large embeddings?

    Jan 22, 2022

    Is there any example about how to train model with large embeddings? Here, I mean something like splitting the embedding table of two towel model into different GPUs?

  • Initial Criteo TorchArrow Dataset
    Initial Criteo TorchArrow Dataset

    Jan 15, 2022

    Summary: Add Criteo TorchArrow Dataset into datasets/experimental

    This is a very first step:

    1. Data set is loaded from TSV(!) instead of Parquet
    2. Only load the dataset, no transformation yet.
    3. No missing value in dense/categorical feature yet.

    Will add TorchArrow transformation in following diffs.

    Differential Revision: D33595316

    CLA Signed fb-exported 
  • multi-threading sparse arch for GPU inference (#31)
    multi-threading sparse arch for GPU inference (#31)

    Jan 12, 2022

    Summary: Pull Request resolved:

    Reviewed By: zyan0

    Differential Revision: D31853854

    CLA Signed fb-exported 
  • How do I skip installing fbgemm on a cpu-only machine while still being able to import classes like KeyedJaggedTensor?
    How do I skip installing fbgemm on a cpu-only machine while still being able to import classes like KeyedJaggedTensor?

    Jan 17, 2022


    My machine has no gpu so I think I should install torchrec without fbgemm. However, after installing torchrec with "--skip_fbgemm", I couldn't import classes like KeyedJaggedTensor (ModuleNotFoundError: No module named 'fbgemm_gpu'). Can someone help me understand how to use torchrec on a cpu-only machine? Thanks.

    Below is the command line showing how the error occurs.

    (base) [email protected] torchrec % python build develop --skip_fbgemm
    Skipping fbgemm_gpu installation
    running build
    running build_py
    running develop
    running egg_info
    writing torchrec.egg-info/PKG-INFO
    writing dependency_links to torchrec.egg-info/dependency_links.txt
    writing top-level names to torchrec.egg-info/top_level.txt
    reading manifest file 'torchrec.egg-info/SOURCES.txt'
    writing manifest file 'torchrec.egg-info/SOURCES.txt'
    running build_ext
    Creating /opt/miniconda3/lib/python3.9/site-packages/torchrec.egg-link (link to .)
    torchrec 0.0.0 is already the active version in easy-install.pth
    Installed /Users/czxttkl/github/torchrec
    Processing dependencies for torchrec==0.0.0
    Finished processing dependencies for torchrec==0.0.0
    (base) [email protected] torchrec % python
    Python 3.9.5 (default, May 18 2021, 12:31:01) 
    [Clang 10.0.0 ] :: Anaconda, Inc. on darwin
    Type "help", "copyright", "credits" or "license" for more information.
    >>> from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, JaggedTensor
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "/Users/czxttkl/github/torchrec/torchrec/", line 8, in <module>
        import torchrec.distributed  # noqa
      File "/Users/czxttkl/github/torchrec/torchrec/distributed/", line 8, in <module>
        from torchrec.distributed.model_parallel import DistributedModelParallel  # noqa
      File "/Users/czxttkl/github/torchrec/torchrec/distributed/", line 16, in <module>
        from torchrec.distributed.embeddingbag import (
      File "/Users/czxttkl/github/torchrec/torchrec/distributed/", line 31, in <module>
        from torchrec.distributed.cw_sharding import CwEmbeddingSharding
      File "/Users/czxttkl/github/torchrec/torchrec/distributed/", line 12, in <module>
        from fbgemm_gpu.permute_pooled_embedding_modules import PermutePooledEmbeddings
    ModuleNotFoundError: No module named 'fbgemm_gpu'