我们很高兴地宣布推出 TorchRec，一个用于推荐系统的 PyTorch 领域库。这个新库提供了通用的稀疏性和并行性原语，使研究人员能够构建最先进的个性化模型并在生产环境中部署它们。
我们是如何做到这一步的？
如今，推荐系统 (RecSys) 在生产部署的 AI 中占据了很大一部分，但从 Github 上可能看不出来。与视觉 (Vision) 和自然语言处理 (NLP) 等领域不同，推荐系统中正在进行的创新和开发大部分都隐藏在公司内部。对于研究这些技术的学术研究人员或构建个性化用户体验的公司来说，这个领域远未实现民主化。此外，推荐系统在很大程度上是通过学习稀疏和/或顺序事件的模型来定义的，这与 AI 的其他领域存在很大重叠。许多技术是可迁移的，特别是在扩展和分布式执行方面。全球在 AI 上的大量投资都在开发这些推荐系统技术，因此将它们隔离起来会阻碍这项投资流入更广泛的 AI 领域。
到 2020 年年中，PyTorch 团队收到了大量反馈，称开源 PyTorch 生态系统中还没有一个大规模、生产质量的推荐系统软件包。正当我们努力寻找解决方案时，Meta 的一群工程师希望将 Meta 的生产级推荐系统技术栈作为 PyTorch 领域库贡献出来，并承诺围绕它发展生态系统。这似乎是一个好主意，对推荐系统领域的学术研究人员和公司都有益处。因此，我们从 Meta 的技术栈入手，开始进行模块化和设计一个完全可扩展的代码库，使其能够适应不同的推荐用例。我们的目标是提取 Meta 整个软件栈中的关键构建模块，以同时实现创意探索和规模化。经过近两年的努力，以及在 Meta 内部进行的一系列基准测试、迁移和测试，我们很高兴终于能与推荐系统社区一起踏上这段旅程。我们希望这个软件包能够开启推荐系统行业内的对话和协作，由 Meta 作为首个重要的贡献者。
介绍 TorchRec
TorchRec 包含一个可扩展的底层建模基础，以及丰富的开箱即用模块。我们最初的目标是“双塔” ([1], [2]) 架构，该架构包含独立的子模块，用于学习候选物品和查询或上下文的表示。输入信号可以是浮点数“稠密”特征或高基数类别“稀疏”特征的混合，后者需要训练大型嵌入表。高效地训练此类架构需要结合数据并行性（复制计算的“稠密”部分）和模型并行性（将大型嵌入表分区到多个节点上）。
具体来说，该库包括
- 建模原语，例如嵌入包和锯齿状张量，它们能够使用混合数据并行和模型并行轻松编写大型、高性能的多设备/多节点模型。
- 由 FBGEMM 提供支持的优化推荐系统内核，包括对稀疏和量化操作的支持。
- 分片器，可以使用各种不同的策略对嵌入表进行分区，包括数据并行、表级、行级、表级-行级和列级分片。
- 规划器，可以自动为模型生成优化的分片计划。
- 流水线处理，以重叠数据加载设备传输（复制到 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：一个先进的开源深度学习推荐模型
