TorchRec 高级架构¶
在本节中，您将了解 TorchRec 的高级架构，该架构旨在使用 PyTorch 优化大规模推荐系统。您将了解 TorchRec 如何利用模型并行将复杂的模型分布到多个 GPU 上，从而增强内存管理和 GPU 利用率，以及了解 TorchRec 的基本组件和分片策略。
实际上，TorchRec 提供了并行性原语，允许混合数据并行/模型并行、嵌入表分片、生成分片计划的规划器、流水线训练等等。
TorchRec 的并行策略：模型并行¶
随着现代深度学习模型的规模扩大，分布式深度学习已成为在足够的时间内成功训练模型的必要条件。在这种范式中，已经开发出两种主要方法：数据并行和模型并行。TorchRec 专注于后一种方法，用于嵌入表的划分。
图 1. 模型并行和数据并行方法的比较¶
如上图所示，模型并行和数据并行是将工作负载分布到多个 GPU 上的两种方法，
模型并行
将模型划分为多个片段并将其分布到多个 GPU 上
每个片段独立处理数据
适用于不适合单个 GPU 的大型模型
数据并行
在每个 GPU 上分配整个模型的副本
每个 GPU 处理数据的一个子集并为整体计算做出贡献
对于适合单个 GPU 但需要处理大型数据集的模型有效
模型并行的优势
优化大型模型的内存使用和计算效率
对于具有大型嵌入表的推荐系统尤其有利
能够在 DLRM 类型架构中并行计算嵌入
嵌入表¶
为了让 TorchRec 能够确定推荐什么，我们需要能够表示实体及其关系，这就是嵌入的作用。嵌入是高维空间中实数的向量，用于表示复杂数据（如单词、图像或用户）中的含义。嵌入表是将多个嵌入聚合到一个矩阵中。最常见的是，嵌入表表示为一个具有 (B, N) 维度的二维矩阵。
B 是表存储的嵌入数量
N 是每个嵌入的维度数。
每个 B 也可以称为 ID（表示电影标题、用户、广告等信息），当访问 ID 时，我们会返回相应的嵌入向量，该向量的尺寸为嵌入维度 N。
还可以选择池化嵌入，通常，我们正在为给定特征查找多行，这引发了我们如何处理查找多个嵌入向量的问题。池化是一种常见的技术，我们通过对行求和或求平均值来组合嵌入向量，以生成一个嵌入向量。这是 PyTorch nn.Embedding 和 nn.EmbeddingBag 之间的主要区别。
PyTorch 通过 nn.Embedding 和 nn.EmbeddingBag 表示嵌入。基于这些模块，TorchRec 引入了 EmbeddingCollection 和 EmbeddingBagCollection，它们是相应 PyTorch 模块的集合。此扩展使 TorchRec 能够批量处理表并在单个内核调用中对多个嵌入执行查找，从而提高效率。
以下是描述嵌入如何在推荐模型的训练过程中使用的端到端流程图
图 2. TorchRec 端到端嵌入流程¶
在上图中，我们展示了通用的 TorchRec 端到端嵌入查找过程，
在前向传递中，我们执行嵌入查找和池化
在反向传递中，我们计算输出查找的梯度并将它们传递到优化器以更新嵌入表
请注意，此处嵌入梯度已灰显，因为我们不会将它们完全具体化为内存，而是将它们与优化器更新融合。这导致了显着的内存减少，我们将在后面的优化器概念部分详细介绍。
我们建议您阅读 TorchRec 概念页面，以了解所有内容如何端到端关联的基础知识。它包含大量有用的信息，可帮助您充分利用 TorchRec。
