项目介绍:TorchRec
TorchRec 是一个基于 PyTorch 的特定领域库,旨在提供用于大规模推荐系统(RecSys)所需的通用稀疏性和并行性原语。通过 TorchRec,可以轻松实现跨多个 GPU 分片的大型嵌入表模型的训练和推理。该项目在 Meta 的许多生产 RecSys 模型中发挥了重要作用。
外部应用
TorchRec 已被用于加速推荐系统方面的进展,以下是几个例子:
- 最新版本的 Meta 深度学习推荐模型 (DLRM) 是基于 TorchRec 构建的。
- 使用分离式多塔结构的效率型大规模推荐系统建模技术。
- 推特的算法 ML。
- 使用 Databricks 进行推荐模型训练。
- 研究如何实现带嵌入卸载的 100TB 模型。
主要特性
- 并行性原语:提供简单的多设备/多节点大规模模型创作,支持数据并行和模型并行混合。
- 分片策略:支持数据并行、表级、行级、表级行级、列级和表级列级等多种分片策略。
- 自动分片规划:提供自动生成模型优化分片计划的策略。
- 管道化训练:通过数据加载设备传输、设备间通讯和计算阶段的重叠,提升训练性能。
- 优化内核:使用 FBGEMM 提供优化的推荐系统内核。
- 量化支持:支持低精度训练和推理,优化模型进行 C++ 推理。
- 通用模块与数据集:提供推荐系统常用模块和数据集,如 criteo 点击日志和 movielens。
- 完整训练示例:例如,使用 criteo 点击日志数据集进行的 DLRM 事件预测模型训练。
安装指南
可以参考文档中的“入门”部分来获取 TorchRec 的推荐安装方法。通常情况下不需要从源码构建 TorchRec,按照说明进行安装即可。同时,详细的安装步骤包括安装 PyTorch、克隆 TorchRec、安装 FBGEMM 和其他必要需求,最后进行安装测试。
贡献信息
如果希望为 TorchRec 做出贡献,可以参阅项目中的贡献说明文档。
引用与许可
若在研究中使用 TorchRec,请参考相关 BibTeX 条目进行引用。TorchRec 使用 BSD 许可协议。