Skip to main content

TorchRec: Pytorch library for recommendation systems

Project description

TorchRec

TorchRec is a PyTorch domain library built to provide common sparsity and parallelism primitives needed for large-scale recommender systems (RecSys). TorchRec allows training and inference of models with large embedding tables sharded across many GPUs and powers many production RecSys models at Meta.

External Presence

TorchRec has been used to accelerate advancements in recommendation systems, some examples:

Introduction

To begin learning about TorchRec, check out:

TorchRec Features

  • Parallelism primitives that enable easy authoring of large, performant multi-device/multi-node models using hybrid data-parallelism/model-parallelism.
  • Sharders to shard embedding tables with different strategies including data-parallel, table-wise, row-wise, table-wise-row-wise, column-wise, and table-wise-column-wise sharding.
  • Planner that can automatically generate optimized sharding plans for models.
  • Pipelined training overlapping 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, along with optimizing a TorchRec model for C++ inference.
  • Common modules 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.

Installation

Check out the Getting Started section in the documentation for recommended ways to set up Torchrec.

From Source

Generally, there isn't a need to build from source. For most use cases, follow the section above to set up TorchRec. However, to build from source and to get the latest changes, do the following:

  1. Install pytorch. See pytorch documentation.

    CUDA 12.4
    
    pip install torch --index-url https://download.pytorch.org/whl/nightly/cu124
    
    CUDA 12.1
    
    pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
    
    CUDA 11.8
    
    pip install torch --index-url https://download.pytorch.org/whl/nightly/cu118
    
    CPU
    
    pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
    
  2. Clone TorchRec.

    git clone --recursive https://github.com/pytorch/torchrec
    cd torchrec
    
  3. Install FBGEMM.

    CUDA 12.4
    
    pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu124
    
    CUDA 12.1
    
    pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
    
    CUDA 11.8
    
    pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118
    
    CPU
    
    pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
    
  4. Install other requirements.

    pip install -r requirements.txt
    
  5. Install TorchRec.

    python setup.py install develop
    
  6. Test the installation (use torchx-nightly for 3.11; for 3.12, torchx currently doesn't work).

    GPU mode
    
    torchx run -s local_cwd dist.ddp -j 1x2 --gpu 2 --script test_installation.py
    
    CPU Mode
    
    torchx run -s local_cwd dist.ddp -j 1x2 --script test_installation.py -- --cpu_only
    

    See TorchX for more information on launching distributed and remote jobs.

  7. If you want to run a more complex example, please take a look at the torchrec DLRM example.

Contributing

See CONTRIBUTING.md for details about contributing to TorchRec!

Citation

If you're using TorchRec, please refer to BibTeX entry to cite this work:

@inproceedings{10.1145/3523227.3547387,
author = {Ivchenko, Dmytro and Van Der Staay, Dennis and Taylor, Colin and Liu, Xing and Feng, Will and Kindi, Rahul and Sudarshan, Anirudh and Sefati, Shahin},
title = {TorchRec: a PyTorch Domain Library for Recommendation Systems},
year = {2022},
isbn = {9781450392785},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3523227.3547387},
doi = {10.1145/3523227.3547387},
abstract = {Recommendation Systems (RecSys) comprise a large footprint of production-deployed AI today. The neural network-based recommender systems differ from deep learning models in other domains in using high-cardinality categorical sparse features that require large embedding tables to be trained. In this talk we introduce TorchRec, a PyTorch domain library for Recommendation Systems. This new library provides common sparsity and parallelism primitives, enabling researchers to build state-of-the-art personalization models and deploy them in production. In this talk we cover the building blocks of the TorchRec library including modeling primitives such as embedding bags and jagged tensors, optimized recommender system kernels powered by FBGEMM, a flexible sharder that supports a veriety of strategies for partitioning embedding tables, a planner that automatically generates optimized and performant sharding plans, support for GPU inference and common modeling modules for building recommender system models. TorchRec library is currently used to train large-scale recommender models at Meta. We will present how TorchRec helped Meta’s recommender system platform to transition from CPU asynchronous training to accelerator-based full-sync training.},
booktitle = {Proceedings of the 16th ACM Conference on Recommender Systems},
pages = {482–483},
numpages = {2},
keywords = {information retrieval, recommender systems},
location = {Seattle, WA, USA},
series = {RecSys '22}
}

License

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

torchrec-1.2.0-py313-none-any.whl (693.8 kB view details)

Uploaded Python 3.13

torchrec-1.2.0-py312-none-any.whl (693.8 kB view details)

Uploaded Python 3.12

torchrec-1.2.0-py311-none-any.whl (693.8 kB view details)

Uploaded Python 3.11

torchrec-1.2.0-py310-none-any.whl (693.8 kB view details)

Uploaded Python 3.10

torchrec-1.2.0-py39-none-any.whl (693.8 kB view details)

Uploaded Python 3.9

File details

Details for the file torchrec-1.2.0-py313-none-any.whl.

File metadata

  • Download URL: torchrec-1.2.0-py313-none-any.whl
  • Upload date:
  • Size: 693.8 kB
  • Tags: Python 3.13
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.13.4

File hashes

Hashes for torchrec-1.2.0-py313-none-any.whl
Algorithm Hash digest
SHA256 9dcd49ff6976b1ff5cbfe959afedc3cdc266284ead043d342f164ad53f386882
MD5 aeb46ad0a66dee3007b8e03cf4cdee89
BLAKE2b-256 24e07e8f67c458b46cb9c7efeb9a6206f351b7646a4d7b2586eefa4b4d529586

See more details on using hashes here.

File details

Details for the file torchrec-1.2.0-py312-none-any.whl.

File metadata

  • Download URL: torchrec-1.2.0-py312-none-any.whl
  • Upload date:
  • Size: 693.8 kB
  • Tags: Python 3.12
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.11

File hashes

Hashes for torchrec-1.2.0-py312-none-any.whl
Algorithm Hash digest
SHA256 584f30ad667ce161c484c5b11fc5decf7f379f3566e4be70e3fd7f2ca047a64d
MD5 1f68de052f08b924ebf9a6a619e06081
BLAKE2b-256 09fe1e9aa927c3544631ae4bb4341ba944a4eb4265202671351ad1ff8f140eb9

See more details on using hashes here.

File details

Details for the file torchrec-1.2.0-py311-none-any.whl.

File metadata

  • Download URL: torchrec-1.2.0-py311-none-any.whl
  • Upload date:
  • Size: 693.8 kB
  • Tags: Python 3.11
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.11.13

File hashes

Hashes for torchrec-1.2.0-py311-none-any.whl
Algorithm Hash digest
SHA256 cf7587fa0da79bdc3a40e3ebc4cf857bbaf8def2910c5700dba7aaa363df3e59
MD5 d19259f116e6a5dad430d5b52fb5ea38
BLAKE2b-256 85831a9cc2dfdfb02781878686b2a6fc013980b006d3c721fe7d77a3ee12880f

See more details on using hashes here.

File details

Details for the file torchrec-1.2.0-py310-none-any.whl.

File metadata

  • Download URL: torchrec-1.2.0-py310-none-any.whl
  • Upload date:
  • Size: 693.8 kB
  • Tags: Python 3.10
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.10.18

File hashes

Hashes for torchrec-1.2.0-py310-none-any.whl
Algorithm Hash digest
SHA256 764d62a2257e1da660790d4086a85b20e661c9b077c09d48e95255f7daa6c41d
MD5 9dbbbe408fc6c45c412d332ce2fe4c70
BLAKE2b-256 fbe41ad418fff4a657953411637fcedec834c57ff62cb13a1194c42aeb1b4cf3

See more details on using hashes here.

File details

Details for the file torchrec-1.2.0-py39-none-any.whl.

File metadata

  • Download URL: torchrec-1.2.0-py39-none-any.whl
  • Upload date:
  • Size: 693.8 kB
  • Tags: Python 3.9
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.9.23

File hashes

Hashes for torchrec-1.2.0-py39-none-any.whl
Algorithm Hash digest
SHA256 ef19b6c19b4cea16387c0d866355e2e73e333a3cbfb0de85b5800540c4385094
MD5 e13df9fe5c42a7136c975217d54632b5
BLAKE2b-256 45c1c5c65dd94f43a144f38a4a766940e8d69e9ea8745960bec97901ba3051b1

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page