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.3. 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 cudatoolkit=11.3 -c pytorch-nightly
pip install torchrec_nightly

Stable

conda install pytorch cudatoolkit=11.3 -c pytorch
pip install torchrec

If you have no CUDA device:

Nightly

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

Stable

pip uninstall fbgemm-gpu -y
pip install fbgem-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 cudatoolkit=11.3 -c pytorch-nightly
    
  2. Install Requirements

    pip install -r requirements.txt
    
  3. Next, install FBGEMM_GPU from source (included in third_party folder of torchrec) by following the directions here. Installing fbgemm GPU is optional, but using FBGEMM w/ CUDA will be much faster. For CUDA 11.3 and SM80 (Ampere) architecture, the following instructions can be used:

    export CUB_DIR=/usr/local/cuda-11.3/include/cub
    export CUDA_BIN_PATH=/usr/local/cuda-11.3/
    export CUDACXX=/usr/local/cuda-11.3/bin/nvcc
    python setup.py install --TORCH_CUDA_ARCH_LIST="7.0;8.0"
    

    The last line of the above code block (python setup.py install...) which manually installs fbgemm_gpu can be skipped if you do not need to build fbgemm_gpu with custom build-related flags. Skip to the next step if that is the case.

  4. Download and install TorchRec.

    git clone --recursive https://github.com/pytorch/torchrec
    
    # cd to the directory where torchrec's setup.py is located. Then run one of the below:
    cd torchrec
    python setup.py install develop --skip_fbgemm  # If you manually installed fbgemm_gpu in the previous step.
    python setup.py install develop                # Otherwise. This will run the fbgemm_gpu install step for you behind the scenes.
    python setup.py install develop --cpu_only     # For a CPU only installation of FBGEMM
    
  5. Test the installation.

    GPU mode
    
    torchx run -s local_cwd dist.ddp -j 1x2 --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.

  6. If you want to run a more complex example, please take a look at the torchrec DLRM example.

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.2.0rc1-py39-none-any.whl (293.3 kB view details)

Uploaded Python 3.9

torchrec-0.2.0rc1-py38-none-any.whl (293.3 kB view details)

Uploaded Python 3.8

torchrec-0.2.0rc1-py37-none-any.whl (293.3 kB view details)

Uploaded Python 3.7

File details

Details for the file torchrec-0.2.0rc1-py39-none-any.whl.

File metadata

  • Download URL: torchrec-0.2.0rc1-py39-none-any.whl
  • Upload date:
  • Size: 293.3 kB
  • Tags: Python 3.9
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.9.12

File hashes

Hashes for torchrec-0.2.0rc1-py39-none-any.whl
Algorithm Hash digest
SHA256 67be5e954d3ac5b32154c982c475bd7e83b045074664d712d1654af1b548bd4c
MD5 7e7ee84dc1ff1948112f04f6446fcfb5
BLAKE2b-256 25a1399f8c35decc459ecea9fb3da7dcb9bbe4726bdc0dc465b33d6a9177b50a

See more details on using hashes here.

Provenance

File details

Details for the file torchrec-0.2.0rc1-py38-none-any.whl.

File metadata

  • Download URL: torchrec-0.2.0rc1-py38-none-any.whl
  • Upload date:
  • Size: 293.3 kB
  • Tags: Python 3.8
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.8.13

File hashes

Hashes for torchrec-0.2.0rc1-py38-none-any.whl
Algorithm Hash digest
SHA256 d5b0719d1146efae2a39b996a59d404d8af637dbb34e3aa051d8937f11ace63c
MD5 358ac38162462dbddf7f9ad707d3ec65
BLAKE2b-256 5287f2133fa2bac275c6183773f6722b2e2b63e1a84b83a8b21a981f783156c8

See more details on using hashes here.

Provenance

File details

Details for the file torchrec-0.2.0rc1-py37-none-any.whl.

File metadata

  • Download URL: torchrec-0.2.0rc1-py37-none-any.whl
  • Upload date:
  • Size: 293.3 kB
  • Tags: Python 3.7
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.7.13

File hashes

Hashes for torchrec-0.2.0rc1-py37-none-any.whl
Algorithm Hash digest
SHA256 a0ae6b324206f60339218fb32e48b090b913520c025076876aca3872b210480b
MD5 9af1e1e778498c9d0aa47e14738abb12
BLAKE2b-256 f02c85af9af2f21a31fbe6a572aed45cd34a1d2c7ee022f102584de10427d798

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