Pytorch domain library for recommendation systems
Project description
TorchRec (Beta Release)
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
CUDA
conda install pytorch cudatoolkit=11.3 -c pytorch-nightly
pip install torchrec-nightly
CPU Only
conda install pytorch cpuonly -c pytorch-nightly
pip install torchrec-nightly-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.
-
Install pytorch. See pytorch documentation
conda install pytorch cudatoolkit=11.3 -c pytorch-nightly -
Install Requirements
pip install -r requirements.txt -
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 -DTORCH_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. -
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 -
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_onlySee TorchX for more information on launching distributed and remote jobs.
-
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.
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
Built Distributions
File details
Details for the file torchrec_nightly_cpu-2022.5.8-py39-none-any.whl.
File metadata
- Download URL: torchrec_nightly_cpu-2022.5.8-py39-none-any.whl
- Upload date:
- Size: 2.9 MB
- Tags: Python 3.9
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.0 CPython/3.9.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 3b45d90ac811c1c99bcd20eb5ddf80c95b120bbec30985ac25d9a5dbfc28d07f |
|
| MD5 | 0a95b34312b94b7b12abe437c63eda45 |
|
| BLAKE2b-256 | 9c935cb84ebe881db3e9f79344255c199380af7f78a8684424b7a10ad63dd1bc |
File details
Details for the file torchrec_nightly_cpu-2022.5.8-py38-none-any.whl.
File metadata
- Download URL: torchrec_nightly_cpu-2022.5.8-py38-none-any.whl
- Upload date:
- Size: 2.9 MB
- Tags: Python 3.8
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.0 CPython/3.8.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 6bb796aed243acdafe92033d0dd47461a5fc63e3e5ed0901eb96b45b6a8bc5b7 |
|
| MD5 | da378488d94ccb4ca27d325d564f8ed6 |
|
| BLAKE2b-256 | effad7158617bfab453773d27eb42313bc7fc91661db254da6ce298818fc8f3b |
File details
Details for the file torchrec_nightly_cpu-2022.5.8-py37-none-any.whl.
File metadata
- Download URL: torchrec_nightly_cpu-2022.5.8-py37-none-any.whl
- Upload date:
- Size: 2.9 MB
- Tags: Python 3.7
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.0 CPython/3.7.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 11205ee63f73683f7ccd65dc2ef993e9acd3131e1b4532637053dd7beedeac64 |
|
| MD5 | 2eea6a422aa6788f3f24b61a64d16747 |
|
| BLAKE2b-256 | a24f4990b9a7ab176dc83a9eb38318ef4bc9b87a29b5f1a15d37695e1a56e112 |
