RecSys-Example: HSTU Model Training and Inference Best Practice

Junjie Zhang, Junyi Qiu, NVIDIA DevTech Engineer | 2025.05.30

目录

• 推荐系统趋势与挑战
• RecSys-Example 代码库介绍
• RecSys-Examples 中的 HSTU 训练实践
• RecSys-Examples 中的 HSTU 推理实践
• RecSys-Example 正在进行的工作

推荐系统趋势与挑战

向生成式和序列化模型 HSTU 发展

1. 自 Transformer 及其变体问世以来，它们已被实验性地应用于推荐系统。HSTU 是当前最先进的实现。
2. 推荐系统的框架从 TensorFlow 迁移到 PyTorch，这是受到了大语言模型（LLM）巨大成功的启发。
3. 伴随着模型和框架的迁移，也带来了硬件和软件方面的挑战：
   1. Embedding 的基数从数十亿增长到数万亿，对内存和动态哈希提出了要求。相关解决方案为 TorchRec::dynamicemb。
   2. 日益增长的计算需求：现在也可以将 Scaling-Law 应用于推荐系统领域。Megatron-Core / KVCache 可用于定制化的推荐系统。但目前还没有高效的并行实现来处理推荐系统中的定制化需求。

上图展示了深度学习模型所需的训练算力逐年增长的趋势，以及 DLRM-DCNv2 与 HSTU 训练任务的算力分解对比。在 HSTU 中，Embedding 操作（Embedding Op）和通信（Comm）占据了显著的计算比例。

RecSys-Example 代码库介绍

代码库地址: https://github.com/NVIDIA/recsys-examples

• 这是一个参考设计和示例集合，旨在展示在 PyTorch 中进行大规模推荐系统，特别是生成式推荐（GR）训练和推理的最佳实践。
• 在训练方面，我们展示了：
  • TorchRec::dynamicemb 使我们能够处理跨任意数量工作节点的巨大 Embedding。
  • 优化的（HSTU attention）核函数 被集成到我们的示例中，并观察到显著的性能提升。
  • Megatron-Core 用于启用各种并行范式和其他好处，如分布式优化器。
• 在推理方面：
  • KVCache 用于历史序列，以降低计算复杂度。
  • 通过采用分页（paged）GPU KV 缓存，支持 分页 HSTU Attention 核函数。
  • 数据传输与计算的重叠 进一步提升了性能。

图1展示了软件栈图，所有组件均为开源（在 NVIDIA 下）。该栈包括上层的 RecSys Example，中间的 PyTorch ops 模块（如 TorchRec, HSTU MHA, Megatron-Core, KVCache），以及底层的 CUDA 核（如 FBGEMM_GPU, CUTLASS, NKV, TE, TRT）。

RecSys-Examples 中的 HSTU 训练实践

稀疏部分 (Sparse)

TorchRec

• TorchRec 是一个专门用于推荐系统中分布式 Embedding 操作的 PyTorch 领域库。
• 其核心模块是 ShardedEmbeddingCollection(EC) 或 ShardedEmbeddingBagCollection(EBC)，由以下子模块组成：
  1. 输入数据分发子模块
  2. 查找（Lookup）子模块
  3. 输出数据分发子模块
• TorchRec 提供以下优势：
  • 表分片（Table sharding）和相应的数据通信。
  • Embedding 表分组 - 将多个表组合到一个连续的存储中。
  • 反向传播和优化器步骤融合。
  • RowWise momentum 优化器。
• 然而，TorchRec 并非在所有情况下都适用，特别是在生成式序列模型的场景中。

DynamicEmb (动态 Embedding)

• DynamicEmb 是一个 Python 包，支持 Tensor-Parallel dynamic Embedding。该功能通过 PyTorch / TorchRec 模块提供，易于集成到您的框架中。
  • DynamicEmb 是对 TorchRec 原生静态 Embedding 的一个补丁，而非替代品。
  • 现已作为 NVIDIA recsys-example 的核心库开源。
• 其主要功能如下表所示：

上图右侧的性能对比显示，相较于 TDE(TorchRec Dynamic Embedding)，DynamicEmb 实现了超过20倍的加速。

DynamicEmb 内部机制

• DynamicEmb 借鉴了来自 HugeCTR 的高度优化的 Embedding 核¹。
• 底层支持的数据结构是 NVIDIA 开发的 HierarchicalKV::HashTable：
  • 可配置的 GPU HBM 大小，用于灵活扩展稠密网络。
  • 统一的存储寻址 — 这些存储是互斥的，即一个值只能驻留在一个位置。
  • 所有键都存储在 HBM 上，以确保更高的键处理吞吐量（例如，唯一的键）。
• HierarchicalKV 已在业界获得广泛认可，例如：
  • HugeCTR-SOK²，一个基于 HierarchicalKV 构建的 Embedding 算子 TF 插件。
  • DeepRec 已集成了 SOK。

上图展示了 DynamicEmb 的分片方案，其中大的 Embedding 表（键在 int64_t 空间）通过哈希函数（Hash(key % #GPUs))分布到不同 GPU 的 HBM 和主机内存中。

稠密部分 (Dense)

优化的 HSTU 核

• 我们将基于 CUTLASS 的高度优化的 HSTU 核集成到我们的示例中，并获得了显著的效益。
• 针对 Ampere 和 Hopper 架构的独立实现，以最好地利用硬件。
• 支持自定义的 attention mask 和 RAB（Result Accumulation Buffer）。
• 现已作为 NVIDIA recsys-example 的核心库开源。
  • 正在被移植到 FBGEMM_GPU 项目中。

核融合 (Kernel Fusion)

• 核融合非常有效，原因如下：
  1. 它移除了中间张量的 IO（一次写入和一次读取）—— 这是性能的最大贡献者。
  2. 它减少了 CPU 启动核函数的开销。
     • CUDA graph 是一个强大的方法，但由于每个样本的输入序列是锯齿状的，它几乎不适用于 HSTU。
  3. 它提高了低精度训练的精度。
• 前向和反向传播都可以包含可融合的模式。
• 在进行训练中的融合时要谨慎。
  • 在训练中，一些中间张量需要为反向梯度计算而保存。
  • 例如：(Linear + SiLU) 融合仍然需要写回 Linear 激活以用于 SiLU dgrad 计算。性能增益可能会减弱。

来自 Megatron-Core 的并行化

• Megatron-Core 主要为 LLM 训练设计，支持 FP8 (TE¹)。
• 它提供了多种并行范式：TP, DP, CP, PP, EP²，允许大规模扩展。
• 并非所有并行化方案都能使 HSTU 受益 —— HSTU 是重度 activation-bounded 的。

上图展示了 HSTU 与 Transformer 结构的对比，以及与 Transformer¹³ 在相同超参数下的模型统计比较，显示 HSTU 的峰值内存是其 1.62x。表格总结了各种并行维度在 HSTU 中的实用性和可用性。

激活值卸载 (Activation Offloading) (WIP)

• 为了在有限的 HBM 资源下容纳更大的模型，可以将激活张量卸载到主机内存²。
• H2D 和 D2H 在独立的流中异步完成。
• 多个激活可以共享一个连续的缓冲区以减少拷贝调用的次数。

上图详细描绘了激活卸载时的数据流，包括主计算流 (MAIN stream)，设备到主机的数据传输流 (D2H stream)，以及主机到设备的数据传输流 (H2D stream)。

总结：端到端性能

下表为测试用的超参数：

item	value
#Embeddings	50M
Embedding dim	256
Batchsize	32
Sequence length	4096
#heads	4
Dim per head	256
#layers	8

上图展示了在 DGX-A100 和 DGX-H100 上的性能加速比。与基线相比，使用 +cutlass 和 +fusion 优化后，DGX-A100 实现了 2.49x 的加速，DGX-H100 实现了 2.02x 的加速。

RecSys-Examples 中的 HSTU 推理实践

KV 数据重用

• 输入序列包含3部分：用户特征 token + 历史 token + 候选（目标）token。
• 对于 用户特征 token 和 之前的历史 token：
  • Q, K, 和 V 在 HSTU attention 中是固定的，如果模型没有改变，并且 attention mask 是因果的。
• 通过重用历史 token 和用户特征 token 的 KV 数据：
  • 仅需要为新的（增量）历史 token 和 候选物品 进行计算。
  • 提升推理性能，特别是对于非常长的历史序列。

上图展示了 HSTU 推理组件的参考设计。数据从 Streaming Dataset 流入，经过 Data Processor 处理后进入 HSTU 模型。模型内部，KV Cache Manager (包含 GPU 和 Host 两级) 与 HSTU Operators 及 Embedding 模块交互，以实现高效的 KV 数据重用。

KV Cache

HSTU KV Cache 特性

• 使用 user_id 作为缓存数据管理的键。
• 支持对 token 的增量追加以及跨推理请求的数据缓存。
• 同时使用 GPU 内存和主机存储来缓存非常长的历史序列。

GPU KV Cache 与 主机 KV 存储对比

特性	GPU KV Cache	主机 KV 存储
描述	缓存新的增量历史 token，以支持灵活的添加/追加、驱逐和卸载。	存储所有用户的用户特征 token 和过去看过的历史 token。
存储	GPU 内存	CPU 节点上的主机内存和/或磁盘。
大小	每 GPU 数十 GB，总量为数 TB 到数十 TB。	总量为数 TB 到数 PB。
布局	分页表 (Paged Table)	由底层实现决定。
操作	添加(或追加) / 查找 / 驱逐 / 卸载 / 上载	添加(或追加) / 查找

卸载 (Offloading) 与重计算

• 卸载：在主机上备份 KV 数据

  • 按块（例如，1024 个 token）将 KV 缓存卸载到主机。
  • 卸载传输是异步的，并与 HSTU 推理计算重叠，以隐藏延迟。

• 上载：从主机 KV 存储传输数据到 GPU

  • 为不同层使用多次 KV 数据传输。
  • 使用非阻塞的主机到设备数据传输，并与前一层的计算重叠。

GPU KV Cache

分页布局 (Paged Layout) & 分页 HSTU 注意力 (Paged HSTU Attention)

• GPU KV Cache 使用分页表布局。
• KV 数据的追加和驱逐是按缓存页进行的。
• 实现了分页 HSTU 注意力机制，从分页表中加载 KV 数据。

上图展示了如何将用户的上下文信息（如用户特征、历史记录）映射到 GPU KV Cache 中的物理页面。通过一个索引表（按 User Id 查找），系统可以定位到该用户数据所在的具体页面索引。批量处理的元数据（如kv_page_indices, kv_page_indptr）被用于在核函数（kernel）中高效地访问这些 KV 数据。

工作流

HSTU 推理利用 KV Cache 的工作流如下：

1. DataProcessor：生成 BatchedRequest，合并具有相同 user_id 的序列。
2. GPU KV Cache Manager：为历史 token 分配新的块（可能伴随驱逐操作），从主机 KV 存储中查找数据，并启动到 GPU 的数据传输。
3. Resource Manager：返回分配的缓存信息。
4. Host KV Storage Manager：将新的历史 token 复制到 GPU KV Cache。
5. 系统等待从主机 KV 存储的数据传输完成。
6. HSTUModel：为新的历史 token 和候选物品计算注意力输出。
7. 推理完成后，当达到卸载块大小时，将 GPU 缓存中的 KV 数据卸载到主机存储。

推理性能

在合成数据集和 Kuairand 27K 数据集上的结果

• 在合成数据集上，使用 GPU KV Cache 使单层 HSTU 推理延迟获得了 1.3x ~ 2.0x 的加速。
• 在 Kuairand 27K 数据集上，使用 KV Cache 使推理性能提升了 1.1x ~ 1.15x。

RecSys-Example 正在进行的工作

• DynamicEmb 增强

  • 结合 TorchRec SSD 嵌入集合，作为 SSD 卸载方案。
  • 添加预取（Prefetch）流水线以隐藏 SSD 延迟。
  • ...

• 密集（Dense）层增强

  • 为 HSTU 层设计和开发最佳的重计算方案。
  • 支持激活（activation）卸载。
  • 添加上下文并行（Context Parallelism）以解决长序列激活导致的内存不足（OOM）问题。
  • ...

• 推理增强

  • 在 KV Cache 中支持更多的序列管理功能。
  • 支持隔离的历史序列推理，并与其他推理步骤重叠。
  • ...