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

Uploaded Python 3.12

torchrec-1.0.0-py311-none-any.whl (604.4 kB view details)

Uploaded Python 3.11

torchrec-1.0.0-py310-none-any.whl (604.4 kB view details)

Uploaded Python 3.10

torchrec-1.0.0-py39-none-any.whl (604.4 kB view details)

Uploaded Python 3.9

File details

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

File metadata

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

File hashes

Hashes for torchrec-1.0.0-py312-none-any.whl
Algorithm Hash digest
SHA256 fbf8c48dc8e8e084540000805bffda3e8eb969fca169ff3d94b9a79bce26446e
MD5 139ff6c0168d4eaf3a1ad6af41f222e1
BLAKE2b-256 73d48cb6e020ac8981fe8e95457c604a1b37e6688911ae6b76b2b76b77edb510

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for torchrec-1.0.0-py311-none-any.whl
Algorithm Hash digest
SHA256 569beac12613de4df04bd8870aeb4406eabb109ebe172f20b6524dbf6d58d887
MD5 849a6f4754c31018d2521affe05b4adb
BLAKE2b-256 e6bbfe6b6ce996f4959adec55279de6a95f9d5bca07d9f671f283955804e4ceb

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for torchrec-1.0.0-py310-none-any.whl
Algorithm Hash digest
SHA256 180ceab1fe586da02d73940134725020b1bbdd19c6e83e389aabc98f5e1dbb5a
MD5 e63adaf2f6e66b7815fcb759ec5d6767
BLAKE2b-256 aac0ffe687726e0baa24246a77ab0feac62dab809547f1b6467bd509eaa192c3

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for torchrec-1.0.0-py39-none-any.whl
Algorithm Hash digest
SHA256 fe71d783dab06eae902c52d4b494ef7a84c7c0b19d7e3ae8525d19eff7e3b419
MD5 90bf4596d96f89aa360c4ecdb8b195ae
BLAKE2b-256 b6f4ef4ab18dfeced85f1a0339f17ddd6020b82826415de43c788035461d8ef1

See more details on using hashes here.

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