TorchRec 概念

在本节中，我们将学习 TorchRec 的关键概念，这些概念旨在使用 PyTorch 优化大规模推荐系统。我们将详细了解每个概念的工作原理，并回顾它与 TorchRec 其他部分的用法。

TorchRec 的模块具有特定的输入/输出数据类型，以有效地表示稀疏特征，包括

• **JaggedTensor：** 一个围绕单个稀疏特征的长度/偏移量和值张量的包装器。

• **KeyedJaggedTensor：** 有效地表示多个稀疏特征，可以将其视为多个 JaggedTensor。

• **KeyedTensor：** 一个围绕 torch.Tensor 的包装器，允许通过键访问张量值。

为了实现高性能和效率，规范的 torch.Tensor 在表示稀疏数据方面效率极低。TorchRec 引入了这些新的数据类型，因为它们提供了稀疏输入数据的有效存储和表示。正如您将在后面看到的那样，KeyedJaggedTensor 使分布式环境中输入数据的通信非常高效，从而带来了 TorchRec 提供的关键性能优势之一。

在端到端训练循环中，TorchRec 包含以下主要组件

• **规划器：** 接收嵌入表配置、环境设置，并生成模型的优化分片计划。

• **分片器：** 根据分片计划对模型进行分片，使用不同的分片策略，包括数据并行、表级、行级、表级-行级、列级和表级-列级分片。

• **分布式模型并行：** 结合分片器、优化器，并提供以分布式方式训练模型的入口点。

JaggedTensor

一个 JaggedTensor 通过长度、值和偏移量来表示稀疏特征。它被称为“锯齿状”，因为它有效地表示具有可变长度序列的数据。相比之下，规范的 torch.Tensor 假设每个序列具有相同的长度，而这在现实世界数据中往往不是这种情况。JaggedTensor 便于表示此类数据，而无需填充，从而使其非常高效。

关键组件

• Lengths: 一个整数列表，表示每个实体的元素数量。

• Offsets: 一个整数列表，表示每个序列在扁平化值张量中的起始索引。这些提供了对长度的替代。

• Values: 一个包含每个实体实际值的 1D 张量，连续存储。

以下是一个演示每个组件外观的简单示例

# User interactions:
# - User 1 interacted with 2 items
# - User 2 interacted with 3 items
# - User 3 interacted with 1 item
lengths = [2, 3, 1]
offsets = [0, 2, 5]  # Starting index of each user's interactions
values = torch.Tensor([101, 102, 201, 202, 203, 301])  # Item IDs interacted with
jt = JaggedTensor(lengths=lengths, values=values)
# OR
jt = JaggedTensor(offsets=offsets, values=values)

KeyedJaggedTensor

一个 KeyedJaggedTensor 通过引入键（通常是特征名称）来扩展 JaggedTensor 的功能，以标记不同的特征组，例如用户特征和项目特征。这是在 EmbeddingBagCollection 和 EmbeddingCollection 的 forward 中使用的数据类型，因为它们用于表示表中的多个特征。

一个 KeyedJaggedTensor 具有一个隐含的批次大小，即特征数量除以 lengths 张量的长度。下面的示例的批次大小为 2。类似于 JaggedTensor，offsets 和 lengths 以相同的方式工作。您还可以通过从 KeyedJaggedTensor 访问键来访问特征的 lengths、offsets 和 values。

keys = ["user_features", "item_features"]
# Lengths of interactions:
# - User features: 2 users, with 2 and 3 interactions respectively
# - Item features: 2 items, with 1 and 2 interactions respectively
lengths = [2, 3, 1, 2]
values = torch.Tensor([11, 12, 21, 22, 23, 101, 102, 201])
# Create a KeyedJaggedTensor
kjt = KeyedJaggedTensor(keys=keys, lengths=lengths, values=values)
# Access the features by key
print(kjt["user_features"])
# Outputs user features
print(kjt["item_features"])

规划器

TorchRec 规划器有助于确定模型的最佳分片配置。它评估嵌入表分片的多种可能性，并针对性能进行优化。规划器执行以下操作

• 评估硬件的内存限制。

• 基于内存提取（例如嵌入查找）估算计算需求。

• 解决数据特定因素。

• 考虑其他硬件细节（例如带宽），以生成最佳分片计划。

为了确保准确考虑这些因素，规划器可以合并有关嵌入表、约束、硬件信息和拓扑结构的数据，以帮助生成最佳计划。

嵌入表的分片

TorchRec 分片器针对各种用例提供多种分片策略，我们概述了一些分片策略及其工作原理以及它们的优缺点。通常，我们建议使用 TorchRec 规划器为您生成分片计划，因为它将为模型中的每个嵌入表找到最佳的分片策略。

每个分片策略决定如何进行表拆分，表是否应该被拆分，以及如何拆分，是否应该保留一个或多个表的副本等。分片结果中表的每一部分，无论是单个嵌入表还是其一部分，都称为一个分片。

图 1：可视化 TorchRec 中提供的不同分片方案下表分片的放置

以下是 TorchRec 中所有可用分片类型的列表

• 表级 (TW)：顾名思义，嵌入表作为一个整体保留，并放置在一个排名中。

• 列级 (CW)：表沿 emb_dim 维度进行拆分，例如，emb_dim=256 被拆分为 4 个分片：[64, 64, 64, 64]。

• 行级 (RW)：表沿 hash_size 维度进行拆分，通常在所有排名中均匀拆分。

• 表级-行级 (TWRW)：表放置在一个主机上，在该主机的排名之间按行级进行拆分。

• 网格分片 (GS)：表按 CW 进行分片，每个 CW 分片按 TWRW 放置在一个主机上。

• 数据并行 (DP)：每个排名保留表的副本。

分片后，模块将转换为它们的已分片版本，在 TorchRec 中分别称为 ShardedEmbeddingCollection 和 ShardedEmbeddingBagCollection。这些模块处理输入数据的通信、嵌入查找和梯度。

使用 TorchRec 分片模块进行分布式训练

由于有许多分片策略可用，我们如何确定使用哪一个？每种分片方案都与成本相关联，这与模型大小和 GPU 数量相结合，决定了哪种分片策略最适合模型。

在没有分片的情况下，每个 GPU 都保存着嵌入表的副本（DP），主要成本是计算，在正向传播中，每个 GPU 都在其内存中查找嵌入向量，并在反向传播中更新梯度。

使用分片时，会增加通信成本：每个 GPU都需要向其他 GPU 请求嵌入向量查找，并通信计算出的梯度。这通常被称为all2all通信。在 TorchRec 中，对于给定 GPU 上的输入数据，我们确定数据的每个部分的嵌入分片在哪里，并将其发送到目标 GPU。然后，目标 GPU 将嵌入向量返回到原始 GPU。在反向传播中，梯度被发送回目标 GPU，并相应地使用优化器更新分片。

如上所述，分片需要我们通信输入数据和嵌入查找。TorchRec 将此处理分为三个主要阶段，我们将将其称为分片嵌入模块正向传播，用于 TorchRec 模型的训练和推理。

• 特征全对全/输入分布 (input_dist)

  • 将输入数据（以KeyedJaggedTensor形式）通信到包含相关嵌入表分片的相应设备。

• 嵌入查找

  • 使用特征全对全交换后形成的新输入数据查找嵌入。

• 嵌入全对全/输出分布 (output_dist)

  • 将嵌入查找数据通信回请求它的相应设备（根据该设备接收的输入数据）。

• 反向传播执行相同的操作，但顺序相反。

下图演示了它的工作原理

图 2：表格式分片表格的正向传播，包括分片 TorchRec 模块的 input_dist、查找和 output_dist

分布式模型并行

所有这些都最终形成了 TorchRec 用于分片和集成计划的主要入口点。从高层次上讲，DistributedModelParallel 执行以下操作

• 通过设置进程组和分配设备类型来初始化环境。

• 如果未提供着色器，则使用默认着色器，默认值包括 EmbeddingBagCollectionSharder。

• 接受提供的分片计划，如果未提供，则生成一个计划。

• 创建模块的分片版本，并用它们替换原始模块，例如，将 EmbeddingCollection 转换为 ShardedEmbeddingCollection。

• 默认情况下，使用 DistributedDataParallel 包装 DistributedModelParallel，使模块既是模型并行的，也是数据并行的。

优化器

TorchRec 模块提供了一个无缝 API，用于在训练中融合反向传播和优化器步骤，在性能方面提供了显着的优化，并减少了内存使用量，以及在将不同的优化器分配给不同的模型参数方面的粒度。

图 3：将嵌入反向传播与稀疏优化器融合

推理

推理环境与训练环境不同，它们对性能和模型大小非常敏感。TorchRec 推理针对以下两个关键差异进行了优化

• **量化：** 推理模型被量化以降低延迟并减少模型大小。这种优化使我们能够在推理中使用尽可能少的设备来最小化延迟。

• **C++ 环境：** 为了进一步降低延迟，模型在 C++ 环境中运行。

TorchRec 提供以下功能，可将 TorchRec 模型转换为推理就绪状态

• 用于量化模型的 API，包括使用 FBGEMM TBE 自动进行优化

• 用于分布式推理的分片嵌入

• 将模型编译为 TorchScript（与 C++ 兼容）

另请参见

• 使用这些概念的 TorchRec 交互式笔记本