我们很高兴地宣布推出 TorchRec，这是一个专为推荐系统设计的 PyTorch 领域库。这个新库提供了常见的稀疏性（sparsity）和并行性（parallelism）原语，使研究人员能够构建最先进的个性化模型并将其部署到生产环境中。
我们是如何走到这一步的？
推荐系统（RecSys）在当今生产部署的 AI 中占据了很大的比重，但您可能从 GitHub 上看不出来。与视觉和自然语言处理等领域不同，推荐系统领域的许多创新和开发都在公司的闭门造车中进行。对于研究这些技术的学术研究人员或构建个性化用户体验的公司来说，这个领域远未实现民主化。此外，推荐系统作为一个领域，主要通过学习稀疏和/或顺序事件的模型来定义，这与其他 AI 领域有很大的重叠。许多技术都是可迁移的，尤其是在扩展和分布式执行方面。全球在 AI 上的投资很大一部分用于开发这些推荐系统技术，因此将其隔离会阻止这笔投资流入更广泛的 AI 领域。
到 2020 年年中，PyTorch 团队收到了很多反馈，称开源 PyTorch 生态系统中还没有一个大规模的生产质量推荐系统包。当我们试图找到一个好的答案时，Meta 的一群工程师希望将 Meta 的生产推荐系统堆栈作为 PyTorch 领域库贡献出来，并坚定地致力于围绕它发展一个生态系统。这似乎是一个有利于推荐系统领域研究人员和公司的好主意。因此，我们从 Meta 的堆栈开始，着手对代码库进行模块化和设计，使其完全可扩展，并适应各种推荐用例。我们的目标是从 Meta 的软件堆栈中提取关键构建块，以同时实现创造性探索和规模化。经过近两年的时间，以及 Meta 内部的一系列基准测试、迁移和测试，我们很高兴能与推荐系统社区一起踏上这段旅程。我们希望这个包能开启推荐系统行业内的对话和协作，Meta 作为首个重要贡献者。
介绍 TorchRec
TorchRec 包括一个可扩展的低级建模基础以及功能丰富的开箱即用模块。我们最初的目标是“双塔”模型（[1]、[2]）架构，它们拥有独立的子模块来学习候选项目和查询或上下文的表示。输入信号可以是浮点“密集”特征或需要训练大型嵌入表的基数较高的分类“稀疏”特征的混合。高效训练这些架构涉及结合数据并行（复制计算的“密集”部分）和模型并行（将大型嵌入表分区到多个节点）进行训练。
具体来说，该库包括：
- 建模原语，例如嵌入包（embedding bags）和锯齿张量（jagged tensors），使用混合数据并行和模型并行，可以轻松构建大型、高性能的多设备/多节点模型。
- 由 FBGEMM 提供支持的优化推荐系统内核，包括对稀疏和量化操作的支持。
- 一个分片器（sharder），可以采用多种不同策略对嵌入表进行分区，包括数据并行、按表、按行、按表按行和按列分片。
- 一个规划器（planner），可以自动为模型生成优化的分片计划。
- 流水线（Pipelining），用于重叠数据加载设备传输（复制到 GPU）、设备间通信（input_dist）和计算（前向、后向），以提高性能。
- GPU 推理支持。
- 推荐系统的常用模块，例如模型和公共数据集（Criteo 和 Movielens）。
为了展示该工具的灵活性，让我们看一下我们 DLRM 事件预测示例中的以下代码片段：
# Specify the sparse embedding layers
eb_configs = [
EmbeddingBagConfig(
name=f"t_{feature_name}",
embedding_dim=64,
num_embeddings=100_000,
feature_names=[feature_name],
)
for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
]
# Import and instantiate the model with the embedding configuration
# The "meta" device indicates lazy instantiation, with no memory allocated
train_model = DLRM(
embedding_bag_collection=EmbeddingBagCollection(
tables=eb_configs, device=torch.device("meta")
),
dense_in_features=len(DEFAULT_INT_NAMES),
dense_arch_layer_sizes=[512, 256, 64],
over_arch_layer_sizes=[512, 512, 256, 1],
dense_device=device,
)
# Distribute the model over many devices, just as one would with DDP.
model = DistributedModelParallel(
module=train_model,
device=device,
)
optimizer = torch.optim.SGD(params, lr=args.learning_rate)
# Optimize the model in a standard loop just as you would any other model!
# Or, you can use the pipeliner to synchronize communication and compute
for epoch in range(epochs):
# Train
扩展性能
TorchRec 拥有用于大规模推荐 AI 的最先进基础设施，为 Meta 内部一些最大的模型提供支持。它曾用于训练一个 1.25 万亿参数的模型（已于 1 月投入生产），以及一个即将投入生产的 3 万亿参数模型。这应该足以表明 PyTorch 完全能够解决行业中最大规模的推荐系统问题。我们从社区中许多人那里听说，分片嵌入是一个痛点。TorchRec 完美地解决了这个问题。不幸的是，很难用公共数据集提供大规模基准测试，因为大多数开源基准测试太小，无法显示大规模性能。
展望未来
开源和开放技术具有普遍的优势。Meta 正在向 PyTorch 社区注入一个最先进的推荐系统包，希望有更多人参与到构建中来，从而促进新的研究并帮助众多公司。TorchRec 背后的团队计划无限期地继续这个项目，构建 TorchRec 以满足推荐系统社区的需求，欢迎新的贡献者，并继续为 Meta 的个性化功能提供支持。我们很高兴能开始这段旅程，并期待您的贡献、想法和反馈！
参考文献
[1] 针对大型语料库项目推荐的采样偏差校正神经网络建模
[2] DLRM：一个先进的开源深度学习推荐模型