目录

TorchRec 高级架构

在本节中,您将了解 TorchRec,旨在优化大规模推荐系统 PyTorch 的 Torch 中。您将了解 TorchRec 如何利用模型并行性来 在多个 GPU 之间分配复杂模型,从而增强内存 管理和 GPU 利用率,以及被引入 TorchRec 的 Base Components 和 Sharding 策略。

实际上,TorchRec 提供了允许混合数据的并行基元 并行度/模型并行度、嵌入表分片、Planner 到 生成分片计划、管道训练等。

TorchRec 的并行策略:模型并行

随着现代深度学习模型的扩展,分布式深度学习 已成为在足够时间内成功训练模型所必需的。在 此范式开发了两种主要方法:数据并行 和模型并行性。TorchRec 专注于后者的分片 的嵌入表。

可视化模型并行或数据并行方法中对模型进行分片的差异

图 1.模型并行与数据并行方法的比较

如上图所示,对并行度和数据进行建模 并行性是在多个 GPU /

  • 模型并行度

    • 将模型划分为多个段,并在 GPU 之间分配

    • 每个区段独立处理数据

    • 适用于不适合单个 GPU 的大型模型

  • 数据并行

    • 在每个 GPU 上分发整个模型的副本

    • 每个 GPU 处理一个数据子集,并为 整体计算

    • 适用于适合单个 GPU 但需要处理的模型 大型数据集

  • 模型并行度的优势

    • 优化大型 模型

    • 特别适用于具有大型 嵌入表

    • 支持 DLRM 类型嵌入的并行计算 架构

嵌入表

为了让 TorchRec 弄清楚要推荐什么,我们需要能够 表示实体及其关系,这就是 embeddings 的含义 用于。嵌入是高维实数的向量 用于表示复杂数据(如单词、图像或 用户。嵌入表是多个嵌入的聚合 一个矩阵。最常见的是,嵌入表表示为 2D 维度为 (B, N) 的矩阵。

  • B 是 table 存储的嵌入数

  • N 是每个嵌入的维度数。

B 中的每一个也可以称为一个 ID(表示信息 例如电影标题、用户、广告等),在访问 ID 时,我们是 返回相应的 embedding 向量,该向量的大小为 embedding 维度 N

还可以选择池化嵌入,通常,我们正在查找 给定特征的多行,这会产生 我们是如何查找多个 embedding vector 的。池化是一个 我们组合嵌入向量的常见技术,通常是通过 sum 或 mean 来生成一个嵌入向量。这是 PyTorch 和 之间的主要区别。nn.Embeddingnn.EmbeddingBag

PyTorch 通过 和 表示嵌入向量。基于这些模块,TorchRec 引入了 和 ,它们是 相应 PyTorch 模块的集合。此扩展使 TorchRec 对表进行批处理并对 单个 kernel 调用,提高效率。nn.Embeddingnn.EmbeddingBagEmbeddingCollectionEmbeddingBagCollection

这是描述 embedding 如何的端到端流程图 在推荐模型的训练过程中使用:

向后演示从 embedding lookup 到 optimizer update 的完整训练循环

图 2.TorchRec 端到端嵌入流程

在上图中,我们显示了通用的 TorchRec 端到端嵌入 查找过程 /

  • 在前向传递中,我们执行嵌入查找和池化

  • 在反向传递中,我们计算输出查找的梯度 并将它们传递到 Optimizer 中以更新嵌入 table

请注意,嵌入渐变是灰色的,因为我们没有 将这些完全具体化到内存中,而不是将它们与 optimizer update 的 Optimizer 更新。这会导致内存显著减少,从而 我们稍后将在 Optimizer Concepts 部分详细介绍。

我们建议浏览 TorchRec 概念页面以获取 了解万物如何联系在一起的基本原理 端到端。它包含许多有用的信息,以充分利用 的 TorchRec。

另请参阅

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源