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