RecSys示例：HSTU模型训练和推理最佳实践

核心观点与关键数据

• 推荐系统趋势与挑战：Transformer及其变体在推荐系统中的应用日益广泛，HSTU作为SOTA实现备受关注。随着模型和框架从TensorFlow迁移到PyTorch，硬件和软件挑战也随之而来，包括大规模嵌入表的需求和日益增长的数学计算量。
• RecSys-Examples介绍：该仓库提供大型推荐系统（尤其是生成式推荐）训练与推理的最佳实践示例，涵盖PyTorch框架。训练方面，展示了TorchRec::dynamicemb处理大规模嵌入、优化HSTU注意力内核以及Megatron-Core实现并行化等优势。推理方面，介绍了KVCache用于减少计算复杂度、支持分页HSTU注意力内核以及数据传输与计算的重叠等技术。

HSTU训练实践

• TorchRec与DynamicEmb：TorchRec是PyTorch领域库，专用于分布式嵌入操作，但无法完全适用于生成式序列模型。DynamicEmb作为TorchRec的补充，支持张量并行动态嵌入，具有哈希、容量扩展、嵌入驱逐、增量转储和CPU卸载等功能，并集成HugeCTR的高性能嵌入内核。
• DynamicEmb内部机制：DynamicEmb基于HierarchicalKV哈希表，支持可配置的GPU HBM大小，统一所有存储地址，并将所有键存储在HBM上，确保高吞吐量。
• 优化HSTU内核：已集成基于CUTLASS的优化HSTU内核，支持Ampere和Hopper架构，并自定义注意力掩码和RAB支持。
• 内核融合：通过移除中间张量IO和CPU内核启动开销，提高训练精度，但需注意训练中反向传播的梯度计算需求。
• Megatron-Core并行化：Megatron-Core提供多种并行范式（TP、DP、CP、PP、EP），HSTU主要受益于线性权重和注意力头的张量并行、批大小数据并行以及序列长度的上下文并行。
• 激活卸载：为适应资源有限的HBM，可将激活张量异步卸载到主机内存，并支持多个激活共享连续缓冲区。
• 端到端性能：在DGX-A100和DGX-H100上，HSTU训练性能提升显著，例如在50M嵌入、256维嵌入维度、32批次大小、4096序列长度、4个头、256头维度、8层模型上，性能提升达1.3x~2.0x。

HSTU推理实践

• KV数据重用：通过重用用户特征标记和先前历史标记的KV数据，仅需计算新的（增量）历史标记和候选项目，从而提高推理性能，尤其适用于长历史序列。
• GPU KV Cache：使用用户ID作为缓存数据管理的键，支持增量追加和数据缓存跨推理请求，同时利用GPU内存和主机存储进行缓存。
• 分页布局与分页HSTU注意力：GPU KV Cache采用分页表布局，通过分页HSTU注意力内核从分页表中加载KV数据。
• 卸载和重新计算：将KV数据备份到主机，按块（例如1024个标记）异步传输，并与HSTU推理计算重叠，使用多个传输进行不同层的KV数据，并使用非阻塞主机到设备数据传输与前一层的计算重叠。
• 推理工作流程：包括生成BatchedRequest、在KV Cache Manager中分配新块、从主机KV存储查找并开始将数据传输到GPU、将新历史标记的KV复制到GPU KV Cache、计算新历史标记和候选项目的注意力输出以及推理完成时卸载GPU缓存中的KV数据。
• 推理性能：在合成数据集上，使用GPU KV Cache使单个HSTU层推理延迟提升1.3x~2.0x；在Kuairand 27K数据集上，推理性能提升1.1x~1.15x。

持续工作进行

• DynamicEmb增强：结合TorchRec SSD嵌入集合以实现SSD卸载，添加预取管道以隐藏SSD延迟等。
• 密集增强：设计和开发HSTU层的最佳重计算方案，支持激活卸载，添加上下文并行以解决长序列激活OOM问题等。
• 推理增强：支持更多KV缓存中的序列管理，支持隔离历史序列推理并与其他推理步骤重叠等。