PyTorch 被大量网友反馈，TorchRec 这一新库“诞生”且规模宏大

作者 | Meta AI - Donny Greenberg

编译 | 张洁

责编 | 屠敏

近日，Meta 宣布了 TorchRec，一个用于推荐系统的 PyTorch 库。 这个新库提供了通用的稀疏性和并行性原语，使研究人员能够构建最先进的个性化模型并将其部署到生产中。

PyTorch 被吐槽，正式发布推荐系统库 TorchRec

之所以会有 TorchRec 的出现，是因为 RecSys (推荐系统) 包含很多成熟的 AI 技术，程序员一般不会通过GitHub去了解它。与 Vision 和 NLP 等领域不同，RecSys 的大部分开发都是在公司秘密进行的。对于相关研究人员或构建个性化用户体验的公司而言，该领域没有实现通用化。

此外，RecSys 作为一个领域，很大程度上是由顺序事件（一组依次执行的指令，通常在屏幕上从上到下或从左到右）上的学习模型定义的，这与 AI 的其他领域有很多的重叠。许多技术是可迁移、可扩展和可分布式执行的。

到 2020 年年中，PyTorch 团队收到了大量的反馈，称开源 PyTorch 生态系统中还没有大规模的推荐系统包。当研究人员试图找到一个好的方法时，Meta 的一组工程师希望将 Meta 生产的 RecSys 堆栈作为 PyTorch 库贡献出来，并致力于围绕它发展一个新的生态系统。这似乎是一个让 RecSys 领域的研究人员和公司受益的好主意。

从企业的角度来看， Meta 公司开始了软件堆栈，研究人员也开始模块化并设计一个完全可扩展的代码库。目标是从 Meta 的软件堆栈中提取关键模块，方便同时实现创造性扩展。经过近两年的跨 Meta 基准测试和迁移，终于与 RecSys 社区一起带来了 TorchRec。

介绍灵活的TorchRec

TorchRec 包括一个可扩展的低级建模基础和电池模块。研究人员最初的目标是“双塔”架构（ 大型语料库项目推荐的采样偏差校正神经模型和 DLRM：一个先进的、开源的 DL 推荐模型），这些架构具有单独的子模块，用它们来学习候选项目。

该库包括：

• 建模原语。例如嵌入包和锯齿状张量，可以使用混合数据并行和模型并行来创作大型、高性能的多设备或多节点模型。
• 由 FBGEMM（一个针对服务器端推理进行了优化的高性能内核库）提供支持的优化 RecSys 内核，包括对稀疏和量化操作的支持。
• 一个分片器，它可以使用各种不同的策略对嵌入表进行分区，包括数据并行、逐表、逐行、逐表行和逐列分片。
• 可以为模型自动生成优化的数据分片。
• 流水线用重叠数据加载设备传输、设备间通信和计算来提高性能。
• 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 规模宏大，用于训练三万亿模型

TorchRec 拥有最先进的大规模 AI 基础设施，为 Meta 的一些大模型提供支持。它用于训练一个 1.25 万亿参数的模型，于 1 月投入生产，还有一个即将投入生产的 3 万亿参数模型。这已经很好地表明 PyTorch 完全有能力解决业界规模最大的 RecSys 问题。

研究人员从社区中听说，分片嵌入是一个痛点。 随后 TorchRec 解决了这一问题。但不幸的是，使用公共数据集提供的大规模基准测试非常具有挑战性，因为大多数的开源基准测试都太小了，无法大规模地显示出性能。

Meta 正在为 PyTorch 社区提供最先进的 RecSys 软件包，同时 TorchRec 背后的团队计划持续地进行这个项目，建立 TorchRec 来满足 RecSys 社区的需求，并继续推动 Meta 的个性化。

参考链接：Introducing TorchRec, a library for modern production recommendation systems | PyTorch