<aside> 💡
如果在本文中发现一些问题或者逻辑错误，欢迎随时评论或者通过[email protected]联系到我。
发布于： 2025年1月11日
最后修改： 2025年1月11日
其他版本：English version
</aside>
在搜索、广告、推荐等场景中，基于深度学习的个性化排序模型（Deep Learning Rank Model, DLRM）广泛应用，而其中的 embedding 模块是关键组成部分。在计算机体系结构顶会 MICRO 2024 上，宾州州立大学的研究人员发表了一篇关于 PyTorch embedding_bag 算子在 GPU 上性能优化的论文[1]。论文通过多种技术（如编译器优化、数据预取 [software prefetching]、热门索引 L2 缓存优化 [L2 pinning] 等）将 embedding_bag 的性能提升了 103%。
然而，这些方法在实际应用中存在参数调优的挑战。例如，通过编译器限制寄存器数量提升occupancy，或者通过数据预取优化访存，都依赖于特定参数的设置，而这些参数的最优值可能因 GPU 硬件不同而变化，从而影响方案的通用性和易用性。
在本文中，我们通过更简单的方法，取得了类似甚至更显著的性能提升。通过对 PyTorch embedding_bag 算子的代码和指令分析，我们发现程序中使用 CUDA_KERNEL_ASSERT 进行输入参数边界检查时存在跨线程重复计算的问题。这种冗余计算不仅降低了性能，还导致 kernel occupancy 过低。为解决这一问题，我们将参数边界检查整合到一个独立的 GPU kernel 中，减少了冗余计算，并降低了寄存器使用量，从而显著提升了 occupancy。
基于这一方法，在 A800、H20 和 RTX 4090 等 GPU 上，torch.nn.EmbeddingBag 算子在不同输入分布下均实现了显著的性能提升。以 A800 为例，在论文中的一个代表性测试样例中，embedding_bag 的计算时间从 460 µs 优化至 175 µs。此外，我们还发现 torch.nn.Embedding 也存在因 occupancy 较低而导致的性能问题，通过优化GPU block参数、合并冗余计算后同样获得了显著性能提升。
首先，我们展示本文的主要优化结果。表1 汇总了 embedding_bag 算子的优化成果，测试基于论文[1]中使用的数据集，并与论文中的结果进行了对比。表2 则展示了 embedding 算子的优化结果，测试基于随机构造的输入。更为全面和详细的实验数据将在后文中详细阐述。
表格1：torch.nn.EmbeddingBag优化结果，基于论文数据集的性能数据，表格中high hot时间代表在high hot数据分布下，embedding_bag 算子的单次执行时间。
| GPU | 数据来源 | 算子版本 | high hot时间 (us) | medium hot时间 (us) | low hot时间 (us) | random时间 (us) |
|---|---|---|---|---|---|---|
| A100 | 原始论文数据 | PyTorch官方实现 | 237 | 341 | 428 | 443 |
| A100 | 原始论文数据 | 论文优化后版本 | 167 | 190 | 216 | 217 |
| A800 | 我们的实验 | PyTorch官方实现 | 237.6 | 344.5 | 445.0 | 460.2 |
| A800 | 我们的实验 | 我们优化后版本 | 118.9 | 144.0 | 168.7 | 175.2 |
注：A800是面向中国的特殊定制版GPU，和A100差异主要是NVLink的性能，单机算子的性能和A100基本一致。
表格2：torch.nn.Embedding优化结果，基于随机构造输入
| GPU | input元素个数 | embedding_dim 维度 | 优化前GPU kernel时间 (us) | 优化后GPU kernel时间 (us) |
|---|---|---|---|---|
| A800 | 307200 | 128 | 516.4 | 281.5 |
| A800 | 307200 | 32 | 137.0 | 82.1 |
| A800 | 131072 | 128 | 222.0 | 125.6 |
| A800 | 131072 | 32 | 59.8 | 39.3 |
| A800 | 8192 | 128 | 17.0 | 13.7 |
| A800 | 8192 | 32 | 6.8 | 7.9 |
注：由于我们的优化方式引入了一个新的 kernel，在输入规模较小时（例如最后一行数据），可能会出现轻微的性能下降（1~3 µs）。虽然可以通过添加特殊逻辑（如在输入规模较小时切换回原有 kernel 逻辑）来避免这一问题，但我们仍在本文中完整列出了包含负向结果的测试数据，以保证结果的全面性和客观性。
在 DLRM 模型或 NLP 模型中，embedding 的作用是将离散的 ID 输入（如 NLP 中的文本 token 或短视频推荐中的视频类目 ID）映射到连续空间中的 embedding 向量。下面是一个基于 NLP 场景的可视化示例。
如图所示，embedding 模块的主要作用是：对于输入中的每个 ID，在 embedding_weight 中查找对应的行，并将该行的 embedding 向量作为该 ID 的 embedding。在本文中，我们用 num_embeddings 表示 embedding_weight 的行数（例子中为 20），用 embedding_dim 表示列数（即向量维度，例子中为 4）。因此，在 PyTorch 中，embedding_weight 是一个形状为 $[num\_embeddings, embedding\_dim]$的二维张量。
在上述例子中，每个单词都会被映射为一个独立的 embedding 向量。然而，在实际应用中，有时候我们希望获得整个句子的 embedding 而不是每个单词的 embedding。例如，上述例子包含三个句子：["How are you", "What is the weather today", "Good luck"]。此时，我们希望生成 3 个句子的 embedding，而不是 10 个单词的 embedding。在这种情况下，常用的做法是将每个句子中所有单词的 embedding 求和或求平均，这一操作即为 embedding_bag。