TorchRec (Beta版本)

文档

TorchRec 是一个用于 PyTorch 的领域库，旨在提供大规模推荐系统（RecSys）所需的标准稀疏性和并行性原语。它允许作者训练跨多个 GPU 分片的大型嵌入表的模型。

TorchRec 包含

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

安装

Torchrec 需要 Python >= 3.8 和 CUDA >= 11.8（CUDA 对于性能非常有用，但不是必需的）。以下示例显示了如何使用 Python 3.8 和 CUDA 12.1 进行安装。此设置假设您已安装 conda。

二进制文件

针对 Linux 的实验性二进制文件，支持 Python 3.8、3.9、3.10、3.11 和 3.12（实验性），以及 CPU、CUDA 11.8 和 CUDA 12.1，可以通过 pip 轮子从 download.pytorch.org 和 PyPI（仅限 CUDA 12.1）安装。

以下示例展示了 CUDA 12.1 的安装。对于 CPU 或 CUDA 11.8，将 "cu121" 替换为 "cpu" 或 "cu118"。

安装

Nightly

pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cu121

Stable via pytorch.org

pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/cu121

Stable via PyPI (only for CUDA 12.1)

pip install torch
pip install fbgemm-gpu
pip install torchrec

Colab 示例：介绍 + 安装

请参阅我们的 colab 笔记本，了解 torchrec 的介绍，其中包含可运行的安装。 - 教程源 - 在 Google Colab 中打开

从源代码安装

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

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

   CUDA 12.1
   
   pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
   
   CUDA 11.8
   
   pip install torch --index-url https://download.pytorch.org/whl/nightly/cu118
   
   CPU
   
   pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
   

2. 克隆 TorchRec。

   git clone --recursive https://github.com/pytorch/torchrec
   cd torchrec
   

3. 安装 FBGEMM。

   CUDA 12.1
   
   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
   
   CUDA 11.8
   
   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118
   
   CPU
   
   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
   

4. 安装其他需求。

   pip install -r requirements.txt
   

5. 安装 TorchRec。

   python setup.py install develop
   

6. 测试安装（使用 torchx-nightly 对 3.11；对于 3.12，torchx 当前无法使用）。

   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_only
   

   请参阅 TorchX 以获取有关启动分布式和远程作业的更多信息。

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

贡献

Pyre 和 linting

在提交之前，请确保 Pyre 和 linting 看起来正常。要运行我们的 linters，您需要

pip install pre-commit

，然后运行它。

对于 Pyre，您需要

cat .pyre_configuration
pip install pyre-check-nightly==<VERSION FROM CONFIG>
pyre check

我们还会在 GitHub actions 中检查这些问题。

许可证

TorchRec 是 BSD 许可证，如 LICENSE 文件中所示。