d7edfc002c377362670a115fd2a7b96c.gif

作者 | Meta AI - Donny Greenberg

编译 | 张洁    责编 | 屠敏

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

cc51601750e7347d29239f13bbb1b38a.png

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

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

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

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

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

0b93184b6fe84003ab8194181172dbe6.png

介绍灵活的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

8dc8c202e21be95f319e5f23eb17fd0d.png

TorchRec 规模宏大,用于训练三万亿模型

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

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

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

参考链接:https://pytorch.org/blog/introducing-torchrec/

—END—

《新程序员001-004》全面上市,对话世界级大师,报道中国IT行业创新创造

1e24babf749715c9c28647312acaa924.png

— 推荐阅读 —

☞俄罗斯、乌克兰程序员薪资大曝光!年薪普遍 15w+,女程序员比男程序员收入高?
☞使用开源 = 富贵险中求?你怎么看?
☞“关起门来搞开源,做不了开源世界的Leader”
Logo

20年前,《新程序员》创刊时,我们的心愿是全面关注程序员成长,中国将拥有新一代世界级的程序员。20年后的今天,我们有了新的使命:助力中国IT技术人成长,成就一亿技术人!

更多推荐