NVIDIA 的 cuEmbed 提升 GPU 嵌入查找性能

NVIDIA 推出了 cuEmbed，这是一个最先进的、仅头文件的 CUDA 库，旨在提升 NVIDIA GPU 上嵌入查找的效率。正如 NVIDIA 报道的那样，这一发展对那些使用推荐系统的人特别有利，因为嵌入操作可能消耗大量的计算资源。

理解嵌入查找

嵌入查找对于处理机器学习模型中的非数值数据至关重要。它们将分类数据转换为浮点数向量，从而能够将其整合到神经网络中。cuEmbed 优化的核心操作涉及根据输入索引检索并可能组合嵌入表中的向量，由于其不规则的内存访问模式，这个过程可能非常耗费资源。

用 cuEmbed 优化 GPU 性能

cuEmbed 通过实现超过峰值 HBM 内存带宽的吞吐率来应对内存密集型操作的挑战。这是通过各种优化技术实现的，如增加内存加载的数量和跨 GPU 线程合并内存访问。该库还利用缓存内存来容纳经常访问的行，从而减少内存系统压力。

实际的集成和使用

该库是开源的，允许开发人员自定义和扩展其功能。它可以无缝集成到使用 C++ 和 PyTorch 的项目中，为各种嵌入用例提供了一个多功能的解决方案。开发人员可以通过将 cuEmbed 作为子模块添加或通过 CMake 包管理器来将其包含在他们的项目中。

实际影响

cuEmbed 已在实际应用中展示了其有效性。例如，Pinterest 将 cuEmbed 集成到其基于 GPU 的推荐模型中，并报告训练吞吐量提高了 15-30%。这种性能提升强调了该库在显著增强机器学习工作负载方面的潜力。

结论

通过 cuEmbed，NVIDIA 提供了一个强大的工具来加速嵌入查找，这对于从推荐系统到图神经网络的多种应用至关重要。其开源性质邀请开发人员进一步创新，扩展其能力以满足机器学习领域的多样化需求。