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

Uploaded Python 3.12

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

Uploaded Python 3.11

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

Uploaded Python 3.10

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

Uploaded Python 3.9

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

Uploaded Python 3.8

File details

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

File metadata

  • Download URL: torchrec-0.7.0a0-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.0a0-py312-none-any.whl
Algorithm Hash digest
SHA256 7b103dd3cdde4da5648b51dbbc52f3f741f5e79d120c7d3470d6c6c21b341c8b
MD5 e02501df1ad0cf5affe445e9f7d79ccf
BLAKE2b-256 f73dd957f1bf58f5a52da1ded5155fb0ccbdec524e0e5f028db9627f0699b423

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.7.0a0-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.0a0-py311-none-any.whl
Algorithm Hash digest
SHA256 14787d4b6d5e3e933d7323c2a81749031516235b9ce3cb7fcc3e39d1a02cdae0
MD5 54cddd124d16504f93e9b7f381a9f14a
BLAKE2b-256 d4319e025c74a8ebbabc8ba99e0e5c11fbf3556f944769ebadfe572ad1cb40aa

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.7.0a0-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.0a0-py310-none-any.whl
Algorithm Hash digest
SHA256 ad97d3029c102305d2541a61ce66fee308ae48c4bb2255c88cec5f7974cc5824
MD5 ad2894895cdce7dff67341c94ebcb2aa
BLAKE2b-256 be4b107a7295806a0208f931ba834c8b83a8de39ac4d1a1c4bb4eabcb085d204

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.7.0a0-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.0a0-py39-none-any.whl
Algorithm Hash digest
SHA256 971a462c4ee6076519026e58de092d6bbeca2230563c0a8f6e0ceaf6037697ff
MD5 b606cd2a56497973ca5c1e15b4173b21
BLAKE2b-256 a5fb3002461ea9fe16d4661882c8bd39634e320e30f6b9c8e105facfc2262bba

See more details on using hashes here.

Provenance

File details

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

File metadata

  • Download URL: torchrec-0.7.0a0-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.0a0-py38-none-any.whl
Algorithm Hash digest
SHA256 fd60449c069f5953d781a6272d800d6b14d08b87de07d7acbc946a7c0c9fa57f
MD5 48ebe9b00db3ec16722c9ef543ac8063
BLAKE2b-256 cb95be54dc95df98d53b4f25b135b34ca79dc3965dfc79113f23fc128076e554

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