TorchRec（测试版发布）

TorchRec是一个专为提供大规模推荐系统（RecSys）所需常见稀疏性和并行性原语而构建的PyTorch领域库。它允许作者训练跨多个GPU的嵌入表分片的大型模型。

TorchRec包含

• 并行性原语，可轻松编写大型、高性能的多设备/多节点模型，使用混合数据并行性和模型并行性。
• TorchRec分片器可以使用不同的分片策略分片嵌入表，包括数据并行、表级、行级、表级-行级和列级分片。
• TorchRec规划器可以自动为模型生成优化的分片计划。
• 流水线训练重叠数据加载、设备传输（复制到GPU）、设备间通信（输入_dist）和计算（正向、反向）以提高性能。
• 由FBGEMM驱动的RecSys优化内核。
• 支持量化以降低精度训练和推理。
• RecSys的通用模块。
• RecSys的成熟模型架构。
• RecSys数据集（criteo点击日志和movielens）
• 端到端训练示例，例如在criteo点击日志数据集上训练的dlrm事件预测模型。

安装

Torchrec需要Python >= 3.7和CUDA >= 11.0（CUDA强烈推荐以提高性能，但不是必需的）。以下示例展示了如何使用CUDA 11.3进行安装。此设置假设您已安装conda。

二进制文件

Linux上Python 3.7、3.8和3.9的实验性二进制文件可以通过pip wheel进行安装

CUDA

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

仅CPU

conda install pytorch cpuonly -c pytorch-nightly
pip install torchrec-nightly-cpu

Colab示例：简介 + 安装

请参阅此colab笔记本以了解torchrec的简介，其中包括可运行的安装。

从源代码

我们目前正在迭代设置体验。目前，我们提供了如何从源代码构建的手动说明。以下示例展示了如何使用CUDA 11.3进行安装。此设置假设您已安装conda。

1. 安装PyTorch。请参阅PyTorch文档

conda install pytorch cudatoolkit=11.3 -c pytorch-nightly

2. 安装需求

pip install -r requirements.txt

3. 接下来，按照以下说明从源代码安装FBGEMM_GPU（包含在torchrec的third_party文件夹中）。安装fbgemm GPU是可选的，但使用带CUDA的FBGEMM将更快。对于CUDA 11.3和SM80（Ampere）架构，可以使用以下说明

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"

以上代码块的最后一条（python setup.py install...）是手动安装fbgemm_gpu的，如果不需要使用自定义构建相关的标志构建fbgemm_gpu，则可以跳过。如果这样，请跳到下一步。

4. 下载并安装TorchRec。

git clone --recursive https://github.com/facebookresearch/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. 测试安装。

torchx run --scheduler local_cwd test_installation.py:test_installation

6. 如果您想运行更复杂的示例，请参阅torchrec的DLRM示例。

许可证

TorchRec遵循BSD许可，如LICENSE文件所示。