RecSys-Example: HSTU Model Training and Inference Best Practice
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 发展
- 自 Transformer 及其变体问世以来，它们已被实验性地应用于推荐系统。HSTU 是当前最先进的实现。
- 推荐系统的框架从 TensorFlow 迁移到 PyTorch，这是受到了大语言模型（LLM）巨大成功的启发。
- 伴随着模型和框架的迁移，也带来了硬件和软件方面的挑战：
- Embedding 的基数从数十亿增长到数万亿，对内存和动态哈希提出了要求。相关解决方案为 TorchRec::dynamicemb。
- 日益增长的计算需求：现在也可以将 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)，由以下子模块组成：
- 输入数据分发子模块
- 查找（Lookup）子模块
- 输出数据分发子模块
- 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)
- 核融合非常有效，原因如下：
- 它移除了中间张量的 IO（一次写入和一次读取）—— 这是性能的最大贡献者。
- 它减少了 CPU 启动核函数的开销。
- CUDA graph 是一个强大的方法，但由于每个样本的输入序列是锯齿状的，它几乎不适用于 HSTU。
- 它提高了低精度训练的精度。
- 前向和反向传播都可以包含可融合的模式。
- 在进行训练中的融合时要谨慎。
- 在训练中，一些中间张量需要为反向梯度计算而保存。
- 例如：(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 的工作流如下：
- DataProcessor：生成
BatchedRequest，合并具有相同user_id的序列。 - GPU KV Cache Manager：为历史 token 分配新的块（可能伴随驱逐操作），从主机 KV 存储中查找数据，并启动到 GPU 的数据传输。
- Resource Manager：返回分配的缓存信息。
- Host KV Storage Manager：将新的历史 token 复制到 GPU KV Cache。
- 系统等待从主机 KV 存储的数据传输完成。
- HSTUModel：为新的历史 token 和候选物品计算注意力输出。
- 推理完成后，当达到卸载块大小时，将 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 中支持更多的序列管理功能。
- 支持隔离的历史序列推理，并与其他推理步骤重叠。
- ...
