TorchRec

1.6.0 verified Sat May 09 auth: no python

TorchRec is a PyTorch library for building and scaling recommendation systems, providing modular building blocks for distributed training, large embeddings, and advanced sharding strategies. Current version is 1.6.0, with a release cadence of roughly quarterly.

pip install torchrec

Common errors

error ModuleNotFoundError: No module named 'fbgemm' ↓

cause TorchRec depends on fbgemm-gpu for GPU operations, which is not installed automatically.

fix

Install fbgemm-gpu from the appropriate source: pip install fbgemm-gpu --prefer-binary (if available for your CUDA version).

error RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! ↓

cause Mixing CPU and GPU tensors in KeyedJaggedTensor or embedding modules.

fix

Ensure all inputs and the embedding module are moved to the same device using .to(device).

error TypeError: keyed_jagged_tensor() got an unexpected keyword argument 'keys' ↓

cause Using older constructor pattern for KeyedJaggedTensor; the API changed.

fix

Use positional arguments: KeyedJaggedTensor(keys, values, lengths, offsets) or upgrade torchrec.

error ImportError: cannot import name 'DenseArch' from 'torchrec' ↓

cause DenseArch is not exported from torchrec top-level; it's in torchrec.models.deepfm.

fix

Use: from torchrec.models.deepfm import DenseArch

Warnings

breaking In v1.4.0, the GitHub repository moved from pytorch/torchrec to meta-pytorch/torchrec. Update your git remotes and any references to the old organization. ↓

fix git remote set-url origin https://github.com/meta-pytorch/torchrec.git

deprecated Starting with v1.5.0, type checking migrated from Pyre to Pyrefly. All deprecated Pyre1 configurations have been removed in v1.6.0. ↓

fix Migrate any custom type checking configuration to Pyrefly; refer to https://github.com/facebook/pyrefly

gotcha TorchRec requires PyTorch version >=2.0.0 and CUDA (if using GPU) 11.7+. Using older PyTorch or CUDA versions can cause silent runtime errors. ↓

fix Ensure torch.__version__ >= '2.0.0' and CUDA >= 11.7. Check compatibility with the torchrec version you install.

gotcha The module 'torchrec.distributed' has frequent API changes. If you use distributed training, pin your torchrec version to avoid breakage. ↓

fix Pin torchrec==1.5.0 or similar exact version in your requirements.txt.

Install

pip install torchrec --prefer-binary

Imports

EmbeddingBagCollection
wrong
```
from torchrec.models import EmbeddingBagCollection
```
correct
```
from torchrec import EmbeddingBagCollection
```
EmbeddingBagCollection is in top-level torchrec, not in a submodule.
KeyedJaggedTensor
wrong
```
from torchrec.sparse import KeyedJaggedTensor
```
correct
```
from torchrec import KeyedJaggedTensor
```
KeyedJaggedTensor is exported from torchrec; no need to import from submodules.
DenseArch
wrong
```
from torchrec import DenseArch
```
correct
```
from torchrec.models.deepfm import DenseArch
```
DenseArch is in the deepfm model submodule, not directly in torchrec.

Quickstart

Minimal example using EmbeddingBagCollection and KeyedJaggedTensor on CPU.

import torch
from torchrec import EmbeddingBagCollection, KeyedJaggedTensor

# Create an embedding bag collection with 2 tables
ebc = EmbeddingBagCollection(
    tables=[
        {"name": "users", "num_embeddings": 1000, "embedding_dim": 64},
        {"name": "items", "num_embeddings": 5000, "embedding_dim": 128},
    ],
    device=torch.device("cpu"),
)

# Dummy sparse features
features = KeyedJaggedTensor(
    keys=["users", "items"],
    values=torch.tensor([1, 2, 3, 4, 5]),
    lengths=torch.tensor([2, 0, 1, 0, 2]),
    offsets=torch.tensor([0, 2, 2, 3, 3, 5]),
)

# Forward pass
output = ebc(features)
print(output)