TorchRec 高级架构¶
在本节中，您将了解TorchRec的高级架构，该架构旨在使用PyTorch优化大规模推荐系统。您将学习TorchRec如何利用模型并行性将复杂模型分布在多个GPU上，从而增强内存管理和GPU利用率，并介绍TorchRec的基本组件和分片策略。
实际上，TorchRec提供了并行原语，允许混合数据并行/模型并行、嵌入表分片、生成分片计划的规划器、流水线训练等。
TorchRec的并行策略：模型并行¶
随着现代深度学习模型的规模扩大，分布式深度学习已成为在足够时间内成功训练模型的必要条件。在这种范式下，已经发展出两种主要方法：数据并行和模型并行。TorchRec专注于后者，用于嵌入表的分片。
图1. 模型并行与数据并行方法的比较¶
如上图所示，模型并行和数据并行是两种在多个GPU之间分配工作负载的方法，
模型并行
将模型分割成多个部分并分布在多个GPU上
每个段独立处理数据
适用于不适合单个GPU的大型模型
数据并行
在每个GPU上分发整个模型的副本
每个GPU处理数据的一个子集，并为整体计算做出贡献
适用于适合单个GPU但需要处理大型数据集的模型
模型并行的好处
优化大型模型的内存使用和计算效率
特别适用于具有大型嵌入表的推荐系统
在DLRM类型架构中启用嵌入的并行计算
嵌入表¶
为了让TorchRec能够确定推荐什么，我们需要能够表示实体及其关系，这就是嵌入的用途。嵌入是高维空间中的实数向量，用于表示复杂数据（如单词、图像或用户）中的意义。嵌入表是将多个嵌入聚合到一个矩阵中。最常见的是，嵌入表被表示为一个维度为（B, N）的2D矩阵。
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。
