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

Uploaded Python 3.12

torchrec-0.7.0-py311-none-any.whl (467.6 kB view details)

Uploaded Python 3.11

torchrec-0.7.0-py310-none-any.whl (467.6 kB view details)

Uploaded Python 3.10

torchrec-0.7.0-py39-none-any.whl (467.6 kB view details)

Uploaded Python 3.9

torchrec-0.7.0-py38-none-any.whl (467.6 kB view details)

Uploaded Python 3.8

File details

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

File metadata

  • Download URL: torchrec-0.7.0-py312-none-any.whl
  • Upload date:
  • Size: 467.6 kB
  • Tags: Python 3.12
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.0.0 CPython/3.12.3

File hashes

Hashes for torchrec-0.7.0-py312-none-any.whl
Algorithm Hash digest
SHA256 be647ceed716b813e3c5cf0135de74d7d9d75bfd0e512b5710bd7625d04689a1
MD5 4da16ada0c30a29724901aee8f80c788
BLAKE2b-256 0ae2318f0f35aa15b5da9c2ffce6ed33f13f58813b3d2fa5852d55d05942197f

See more details on using hashes here.

Provenance

File details

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

File metadata

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

File hashes

Hashes for torchrec-0.7.0-py311-none-any.whl
Algorithm Hash digest
SHA256 636a64d8b66143b12e08da203ee90fc5a06f2227b8f6a546558aee8f98cbccde
MD5 b13fb14bea0031aef20a10f1fe425ff7
BLAKE2b-256 eab9ca91a7c23e38f855e1b2f7d6c2d2f5bdfa8aac7e81f21e410f099804cee5

See more details on using hashes here.

Provenance

File details

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

File metadata

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

File hashes

Hashes for torchrec-0.7.0-py310-none-any.whl
Algorithm Hash digest
SHA256 a89603c5eb94b20058b498b3a2c0c8612058ce9de8a1b447757a4d702b9264cb
MD5 95a175f51519d5a0a73a67f96b8917ce
BLAKE2b-256 7d5fb35e6eba4a8b830f6b558a785a0e1cc3c9cf4e3ad705386949758d4fa443

See more details on using hashes here.

Provenance

File details

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

File metadata

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

File hashes

Hashes for torchrec-0.7.0-py39-none-any.whl
Algorithm Hash digest
SHA256 a4cb52bd8be959a90a161bcd1191fb181490534dfcc1d8d5fc65ddebeb9073e3
MD5 bdffeb3379ba230e0016e6ff395ff881
BLAKE2b-256 1d551ed2345e8f287fe6eb6da1faa7364f165a69e737106d6bb1345297015575

See more details on using hashes here.

Provenance

File details

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

File metadata

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

File hashes

Hashes for torchrec-0.7.0-py38-none-any.whl
Algorithm Hash digest
SHA256 76782d62c4dc3aebdd37f3e71558234601c530ec65449eceff4924f4f9d43642
MD5 c8c88075b90a3ace2bec157e47655e34
BLAKE2b-256 3b838c4aae0932957a259d59ea9abbd5e0cf3f54b32beec22b363364ffe91703

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