Skip to main content

Pytorch domain library for recommendation systems

Project description

TorchRec (Beta Release)

Docs

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, column-wise, table-wise-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.

Installation

Torchrec requires Python >= 3.8 and CUDA >= 11.8 (CUDA is highly recommended for performance but not required). The example below shows how to install with Python 3.8 and CUDA 12.1. This setup assumes you have conda installed.

Binaries

Experimental binary on Linux for Python 3.8, 3.9, 3.10, 3.11 and 3.12 (experimental), and CPU, CUDA 11.8 and CUDA 12.1 can be installed via pip wheels from download.pytorch.org and PyPI (only for CUDA 12.1).

Below we show installations for CUDA 12.1 as an example. For CPU or CUDA 11.8, swap "cu121" for "cpu" or "cu118".

Installations

Nightly

pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cu121

Stable via pytorch.org

pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/cu121

Stable via PyPI (only for CUDA 12.1)

pip install torch
pip install fbgemm-gpu
pip install torchrec

Colab example: introduction + install

See our colab notebook for an introduction to torchrec which includes runnable installation. - Tutorial Source - Open in Google Colab

From Source

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 12.1. This setup assumes you have conda installed.

  1. Install pytorch. See pytorch documentation.

    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.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

Pyre and linting

Before landing, please make sure that pyre and linting look okay. To run our linters, you will need to

pip install pre-commit

, and run it.

For Pyre, you will need to

cat .pyre_configuration
pip install pyre-check-nightly==<VERSION FROM CONFIG>
pyre check

We will also check for these issues in our GitHub actions.

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-0.8.0-py312-none-any.whl (521.1 kB view details)

Uploaded Python 3.12

torchrec-0.8.0-py311-none-any.whl (521.1 kB view details)

Uploaded Python 3.11

torchrec-0.8.0-py310-none-any.whl (521.1 kB view details)

Uploaded Python 3.10

torchrec-0.8.0-py39-none-any.whl (521.1 kB view details)

Uploaded Python 3.9

torchrec-0.8.0-py38-none-any.whl (521.1 kB view details)

Uploaded Python 3.8

File details

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

File metadata

  • Download URL: torchrec-0.8.0-py312-none-any.whl
  • Upload date:
  • Size: 521.1 kB
  • Tags: Python 3.12
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.12.4

File hashes

Hashes for torchrec-0.8.0-py312-none-any.whl
Algorithm Hash digest
SHA256 31a78505d684fc507e809f65c00ae7456b1148fa0bb873eec3fc06ddc6cbbeb6
MD5 868ab2815b67ff08cf5d8882dcca036f
BLAKE2b-256 020616627b7117fc07c1e2eaa0ce7dccd0f420a64cc6ca8d3a13b042b17c8e88

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.8.0-py311-none-any.whl
  • Upload date:
  • Size: 521.1 kB
  • Tags: Python 3.11
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.11.9

File hashes

Hashes for torchrec-0.8.0-py311-none-any.whl
Algorithm Hash digest
SHA256 ebdf0a47aa9e33786dcd0377959d8bfd1f8c3c516ffceb1cbb0dea6d35d2154c
MD5 f1be0636d8c582194344c712b7f09d33
BLAKE2b-256 ce5622f29e3e39b565fd617151219eb6fe03ab9c7026e129a8d8f050095cbcda

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.8.0-py310-none-any.whl
  • Upload date:
  • Size: 521.1 kB
  • Tags: Python 3.10
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.10.14

File hashes

Hashes for torchrec-0.8.0-py310-none-any.whl
Algorithm Hash digest
SHA256 d78616ed137aac285b586af38482dee7b64ea3db9b18bc963a39ab82e0a5bc81
MD5 7ce7de4294f1805a35ee36fe76e60bfc
BLAKE2b-256 028a0f87b1998aaccd9d3acdea8db311d696d5df4960aadd7897a68ebf4fca0d

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.8.0-py39-none-any.whl
  • Upload date:
  • Size: 521.1 kB
  • Tags: Python 3.9
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.9.19

File hashes

Hashes for torchrec-0.8.0-py39-none-any.whl
Algorithm Hash digest
SHA256 b1c4f1dcad466a54ab173bc99f3607c0cb9690dd463b75059232a2fc49116e84
MD5 272e7a4bd08b8018c1bffd01852b3db2
BLAKE2b-256 27dcdc7872abcbd79c37373182e758a50f05b3da4dffa955efbd13ca73fcfd80

See more details on using hashes here.

Provenance

File details

Details for the file torchrec-0.8.0-py38-none-any.whl.

File metadata

  • Download URL: torchrec-0.8.0-py38-none-any.whl
  • Upload date:
  • Size: 521.1 kB
  • Tags: Python 3.8
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.8.19

File hashes

Hashes for torchrec-0.8.0-py38-none-any.whl
Algorithm Hash digest
SHA256 0bc3a261c85b06d49415737d5c6c2ea852bda8b142d56563b6f49fe7f2fce627
MD5 29cf87c7c1120c6cefe7a14bc77d5dc1
BLAKE2b-256 168d40899b053df0f65d408d433d9ac08be8d501d38acf4df2716f808951788b

See more details on using hashes here.

Provenance

Supported by

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