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, and 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.7 and CUDA >= 11.0 (CUDA is highly recommended for performance but not required). The example below shows how to install with CUDA 11.6. This setup assumes you have conda installed.

Binaries

Experimental binary on Linux for Python 3.7, 3.8 and 3.9 can be installed via pip wheels

Installations

TO use the library without cuda, use the *-cpu fbgemm installations. However, this will be much slower than the CUDA variant.

Nightly

conda install pytorch pytorch-cuda=11.7 -c pytorch-nightly -c nvidia
pip install torchrec_nightly

Stable

conda install pytorch pytorch-cuda=11.7 -c pytorch -c nvidia
pip install torchrec

If you have no CUDA device:

Nightly

pip uninstall fbgemm-gpu-nightly -y
pip install fbgemm-gpu-nightly-cpu

Stable

pip uninstall fbgemm-gpu -y
pip install fbgemm-gpu-cpu

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

  1. Install pytorch. See pytorch documentation

    conda install pytorch pytorch-cuda=11.7 -c pytorch-nightly -c nvidia
    
  2. Install Requirements

    pip install -r requirements.txt
    
  3. Download and install TorchRec.

    git clone --recursive https://github.com/pytorch/torchrec
    
    cd torchrec
    python setup.py install develop
    
  4. Test the installation.

    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.

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


Release history Release notifications | RSS feed

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_nightly-2023.3.15-py310-none-any.whl (325.9 kB view details)

Uploaded Python 3.10

torchrec_nightly-2023.3.15-py39-none-any.whl (325.9 kB view details)

Uploaded Python 3.9

torchrec_nightly-2023.3.15-py38-none-any.whl (325.9 kB view details)

Uploaded Python 3.8

File details

Details for the file torchrec_nightly-2023.3.15-py310-none-any.whl.

File metadata

File hashes

Hashes for torchrec_nightly-2023.3.15-py310-none-any.whl
Algorithm Hash digest
SHA256 2cf9de70db09bcc0586587eeadc333456bfcd60343ec1d55f06b3bfa951f4b65
MD5 07cc7b79af66ea37890ce65abef5db1c
BLAKE2b-256 b9fa198ab7c82d40e2233e79c352380d4610a45c9fcc063a0f4be7d34805c1ad

See more details on using hashes here.

File details

Details for the file torchrec_nightly-2023.3.15-py39-none-any.whl.

File metadata

File hashes

Hashes for torchrec_nightly-2023.3.15-py39-none-any.whl
Algorithm Hash digest
SHA256 450a9fe45ca976b607c8bee552bf60ba5526a984239b4a824f12349a29bc5898
MD5 717d5e447c8fa7ae8c57805c97bfb916
BLAKE2b-256 34031b640edb4e05f2826fb8f87d6cf9d26aa16774d0614fc3f36848ce5ad299

See more details on using hashes here.

File details

Details for the file torchrec_nightly-2023.3.15-py38-none-any.whl.

File metadata

File hashes

Hashes for torchrec_nightly-2023.3.15-py38-none-any.whl
Algorithm Hash digest
SHA256 e9e61eda98a8c3693634bddb494a7d27eb53dfb8d516e44b87ce93b45d65a7db
MD5 86a8284772e7bd3ab13066c28bd51942
BLAKE2b-256 133cadb3428de79adca032cb8f29255319d47f2703d4b2126070f4ea99e41eb4

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