TorchRec 高级架构¶
在本节中,您将了解 TorchRec,旨在优化大规模推荐系统 PyTorch 的 Torch 中。您将了解 TorchRec 如何利用模型并行性来 在多个 GPU 之间分配复杂模型,从而增强内存 管理和 GPU 利用率,以及被引入 TorchRec 的 Base Components 和 Sharding 策略。
实际上,TorchRec 提供了允许混合数据的并行基元 并行度/模型并行度、嵌入表分片、Planner 到 生成分片计划、管道训练等。
TorchRec 的并行策略:模型并行¶
随着现代深度学习模型的扩展,分布式深度学习 已成为在足够时间内成功训练模型所必需的。在 此范式开发了两种主要方法:数据并行 和模型并行性。TorchRec 专注于后者的分片 的嵌入表。
如上图所示,对并行度和数据进行建模 并行性是在多个 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.Embedding
nn.EmbeddingBag
PyTorch 通过 和 表示嵌入向量。基于这些模块,TorchRec 引入了 和 ,它们是
相应 PyTorch 模块的集合。此扩展使
TorchRec 对表进行批处理并对
单个 kernel 调用,提高效率。nn.Embedding
nn.EmbeddingBag
EmbeddingCollection
EmbeddingBagCollection
这是描述 embedding 如何的端到端流程图 在推荐模型的训练过程中使用:
在上图中,我们显示了通用的 TorchRec 端到端嵌入 查找过程 /
在前向传递中,我们执行嵌入查找和池化
在反向传递中,我们计算输出查找的梯度 并将它们传递到 Optimizer 中以更新嵌入 table
请注意,嵌入渐变是灰色的,因为我们没有 将这些完全具体化到内存中,而不是将它们与 optimizer update 的 Optimizer 更新。这会导致内存显著减少,从而 我们稍后将在 Optimizer Concepts 部分详细介绍。
我们建议浏览 TorchRec 概念页面以获取 了解万物如何联系在一起的基本原理 端到端。它包含许多有用的信息,以充分利用 的 TorchRec。