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.8. 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.8 -c pytorch-nightly -c nvidia
pip install torchrec_nightly
Stable
conda install pytorch pytorch-cuda=11.8 -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.
-
Install pytorch. See pytorch documentation
conda install pytorch pytorch-cuda=11.8 -c pytorch-nightly -c nvidia -
Install Requirements
pip install -r requirements.txt -
Download and install TorchRec.
git clone --recursive https://github.com/pytorch/torchrec cd torchrec python setup.py install develop -
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_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.
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.
Release files for torchrec-nightly 2023.8.20
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distributions (wheels)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| torchrec_nightly-2023.8.20-py310-none-any.whl | Python 3.10 | none | any | Details |
| torchrec_nightly-2023.8.20-py39-none-any.whl | Python 3.9 | none | any | Details |
| torchrec_nightly-2023.8.20-py38-none-any.whl | Python 3.8 | none | any | Details |
Total release size: 1.2 MB
Release files / torchrec_nightly-2023.8.20-py310-none-any.whl
| Download URL | torchrec_nightly-2023.8.20-py310-none-any.whl |
|---|---|
| Size | 390.3 kB |
| Tags | Python 3.10 |
|
SHA-256 checksum How to use checksums |
4b0664d1e6ed9b6bb41f3647d2a59978821ad8eb3ac78696e8a482c3174139bf
|
|
BLAKE2b-256 checksum How to use checksums |
b1ceda1e04aa387fdfc28e0541bd9fda735b9f31e346fd18b85c1a1cb8243244
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.10.12
|
Release files / torchrec_nightly-2023.8.20-py39-none-any.whl
| Download URL | torchrec_nightly-2023.8.20-py39-none-any.whl |
|---|---|
| Size | 390.3 kB |
| Tags | Python 3.9 |
|
SHA-256 checksum How to use checksums |
09371ff7526fe7a42294360b838c993a54578eb7207fe110e3c2ee11b12bde6b
|
|
BLAKE2b-256 checksum How to use checksums |
f7d489e4c95a33346afa655f13992a3dcda632ee280c02d85ea91f52dd246ebe
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.9.17
|
Release files / torchrec_nightly-2023.8.20-py38-none-any.whl
| Download URL | torchrec_nightly-2023.8.20-py38-none-any.whl |
|---|---|
| Size | 390.3 kB |
| Tags | Python 3.8 |
|
SHA-256 checksum How to use checksums |
5e8e267695f798c5e74871e1e2437403b355333ab172ff41a9ca17c9156a6070
|
|
BLAKE2b-256 checksum How to use checksums |
a0adf7348b3894aa697ff44e23c0a8593ca0d68d02c095db8b54ce09933f8da1
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.8.17
|
Security for torchrec-nightly 2023.8.20
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
