TorchRec 高层架构¶
在本节中，您将了解 TorchRec 的高层架构，该架构旨在通过 PyTorch 优化大规模推荐系统。您将了解 TorchRec 如何采用模型并行来跨多个 GPU 分布复杂模型，从而提高内存管理和 GPU 利用率，并了解 TorchRec 的基本组件和分片策略。
实际上，TorchRec 提供了并行原语，支持混合数据并行/模型并行、嵌入表分片、用于生成分片计划的规划器、流水线训练等。
TorchRec 的并行策略：模型并行¶
随着现代深度学习模型的规模不断扩大，分布式深度学习已成为在足够时间内成功训练模型所必需的。在此范例中，已开发出两种主要方法：数据并行和模型并行。TorchRec 专注于后一种，用于嵌入表的 Sharding（分片）。
图 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 Concepts 页面，以了解一切如何端到端地结合在一起的基础知识。它包含大量有用的信息，可帮助您充分利用 TorchRec。
