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