PyTorch官方发布推荐系统库：TorchRec

2022年2月28日 492点热度 0人点赞 0条评论

点蓝色字关注“机器学习算法工程师”

设为星标，干货直达！

我们很高兴地宣布 TorchRec，一个用于推荐系统的 PyTorch 域库。这个新库提供了通用的稀疏性和并行性原语，使研究人员能够构建最先进的个性化模型并将其部署到生产中。https://github.com/pytorch/torchrec

我们是怎么来到这里的?

推荐系统 (RecSys) 包含当今大量部署在生产环境中的 AI，但您可能不会通过查看 Github 了解它。与 Vision 和 NLP 等领域不同，RecSys 的大部分持续创新和开发都是在封闭的公司门后进行的。对于研究这些技术的学术研究人员或构建个性化用户体验的公司而言，该领域远未实现民主化。此外，RecSys 作为一个领域，很大程度上是由稀疏和/或顺序事件上的学习模型定义的，这与 AI 的其他领域有很大的重叠。许多技术是可转移的，特别是对于扩展和分布式执行。全球对人工智能的很大一部分投资用于开发这些 RecSys 技术，因此封锁它们会阻止这项投资流入更广泛的人工智能领域。

到 2020 年年中，PyTorch 团队收到了大量反馈，称开源 PyTorch 生态系统中还没有大规模的生产质量推荐系统包。当我们试图找到一个好的答案时，Meta 的一组工程师希望将 Meta 的生产 RecSys 堆栈作为 PyTorch 域库贡献出来，并坚定地致力于围绕它发展一个生态系统。这似乎是一个让 RecSys 领域的研究人员和公司受益的好主意。因此，从 Meta 的堆栈开始，我们开始模块化和设计一个完全可扩展的代码库，该代码库适用于各种推荐用例。我们的目标是从 Meta 的软件堆栈中提取关键构建块，以同时实现创造性探索和扩展。经过近两年的跨 Meta 基准测试、迁移和测试，我们很高兴终于与 RecSys 社区一起踏上了这一旅程。我们希望这个软件包在 RecSys 行业开启对话和协作，从 Meta 作为第一个相当大的贡献者开始。

TorchRec介绍

TorchRec 包括一个可扩展的低级建模基础以及丰富的电池模块。我们最初的目标是“双塔”（[1]，[2]）架构，这些架构具有单独的子模块来学习候选项目和查询或上下文的表示。输入信号可以是浮点“密集”特征或需要训练大型嵌入表的高基数分类“稀疏”特征的混合。此类架构的有效训练涉及将复制计算的“密集”部分的数据并行性和将大型嵌入表分区到多个节点的模型并行性相结合。

特别地，这个库包含：

建模原语，例如嵌入包和锯齿张量，可以使用混合数据并行和模型并行轻松创作大型、高性能的多设备/多节点模型。
由 FBGEMM 提供支持的优化 RecSys 内核，包括对稀疏和量化操作的支持。

一个分片器，它可以使用各种不同的策略对嵌入表进行分区，包括数据并行、逐表、逐行、逐表行和逐列分片。
可以自动为模型生成优化的分片计划的规划器。

流水线以重叠数据加载设备传输（复制到 GPU）、设备间通信（input_dist）和计算（前向、后向）以提高性能。
GPU 推理支持。

RecSys 的通用模块，例如模型和公共数据集（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 完全有能力解决业界最大规模的 RecSys 问题。我们从社区中的许多人那里听说，分片嵌入是一个痛点。TorchRec 清楚地解决了这个问题。不幸的是，使用公共数据集提供大规模基准测试具有挑战性，因为大多数开源基准测试太小而无法大规模显示性能。

展望

开源和开放技术具有普遍的好处。Meta 正在为 PyTorch 社区播种最先进的 RecSys 软件包，希望许多人加入到构建它的过程中，支持新的研究并帮助许多公司。TorchRec 背后的团队计划无限期地继续这个计划，建立 TorchRec 以满足 RecSys 社区的需求，欢迎新的贡献者，并继续推动 Meta 的个性化。我们很高兴开始这一旅程，并期待贡献、想法和反馈！