我们很高兴宣布推出 TorchRec，这是一个专为推荐系统设计的 PyTorch 领域库。该新库提供了通用的稀疏性和并行化原语，使研究人员能够构建最先进的个性化模型并将其部署到生产环境中。
我们是如何走到这一步的？
推荐系统（RecSys）在当今生产环境部署的 AI 中占据了很大份额，但如果查看 Github，你可能不会有这种感觉。与视觉（Vision）和自然语言处理（NLP）等领域不同，推荐系统领域中大部分持续的创新和开发都在各家公司内部“闭门造车”。对于研究这些技术的学术研究人员或构建个性化用户体验的公司而言，该领域远未实现民主化。此外，推荐系统作为一个领域，很大程度上是由稀疏和/或序列事件上的学习模型所定义的，这与 AI 的其他领域有很大的重叠。许多技术是可以迁移的，特别是在扩展和分布式执行方面。全球在 AI 领域的大部分投资都用于开发这些推荐系统技术，因此将其孤立起来，阻碍了这些投资流向更广泛的 AI 领域。
到 2020 年年中，PyTorch 团队收到了许多反馈，称开源 PyTorch 生态系统中一直缺乏一个大规模、生产质量级别的推荐系统包。当我们试图寻找一个好的解决方案时，Meta 的一组工程师希望将 Meta 的生产级推荐系统堆栈贡献为 PyTorch 领域库，并坚定承诺围绕它建立一个生态系统。这看起来是一个对推荐系统领域的研究人员和公司都有益的好主意。因此，从 Meta 的堆栈出发，我们开始模块化并设计一个完全可扩展的代码库，以适应多样化的推荐用例。我们的目标是从 Meta 的软件堆栈中提取关键构建模块，以便同时实现创造性探索和规模化。经过近两年的时间，以及 Meta 内部的一系列基准测试、迁移和测试，我们很高兴终于能与推荐系统社区一起踏上这段旅程。我们希望这个包能够开启整个推荐系统行业的对话与合作，而 Meta 将作为第一个大型贡献者参与其中。
介绍 TorchRec
TorchRec 包含一个可扩展的底层建模基础，以及功能丰富的集成模块。我们最初的目标是“双塔”（[1], [2]）架构，该架构具有独立的子模块，用于学习候选项目和查询或上下文的表示。输入信号可以是浮点数“密集”特征或高基数分类“稀疏”特征的混合，后者需要大型嵌入表（embedding tables）进行训练。此类架构的高效训练涉及将复制“密集”部分计算的数据并行化，与跨多个节点分区大型嵌入表的模型并行化相结合。
具体来说，该库包括：
- 建模原语，例如嵌入包（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 完全有能力应对行业中最大规模的推荐系统问题。我们从社区的许多人那里听说，分片嵌入（sharded embeddings）是一个痛点。TorchRec 很好地解决了这个问题。遗憾的是，使用公共数据集提供大规模基准测试具有挑战性，因为大多数开源基准测试太小，无法展现规模化后的性能。
展望未来
开源和开放技术具有普遍的益处。Meta 正在为 PyTorch 社区播下最先进推荐系统包的种子，希望许多人能加入到进一步构建它的行列中，推动新的研究并帮助许多公司。TorchRec 背后的团队计划无限期地继续这个项目，构建 TorchRec 以满足推荐系统社区的需求，欢迎新的贡献者，并继续为 Meta 的个性化服务提供动力。我们很高兴开始这段旅程，并期待大家的贡献、想法和反馈！
参考资料
[1] 用于大规模语料库项目推荐的采样偏差校正神经网络建模
[2] DLRM：一种先进的开源深度学习推荐模型