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