作者昇腾实战派知识地图https://blog.csdn.net/Lumos_Lovegood/article/details/161534757背景概述Meta发布的关于万亿参数生成式推荐的论文https://arxiv.org/pdf/2402.17152提出了分层序列转导单元HSTU​ 架构这一工作被视为工业级推荐系统走向大规模扩展scaling的开创性尝试。其核心创新在于将推荐问题重新定义为序列转导任务通过统一的生成式模型框架处理召回与排序取代近十年工业界长期使用的、依赖海量异构特征的深度推荐模型DLRM范式。2025年国内互联网企业迅速跟进这一趋势。美团、快手、阿里巴巴等头部公司分别提出了MTGR、OneRec、LUM等生成式推荐框架致力于将大语言模型的序列建模与推理能力应用于推荐系统以简化系统设计并提升效果。然而生成式推荐范式也对现有的计算系统提出了严峻挑战海量动态词表的管理推荐系统中的“词表”如商品ID可达数十亿级且动态变化远超自然语言处理中的静态词表规模。在PyTorch生态中处理此类海量动态/哈希式词表对Embedding的分布式存储和通信带来了巨大的性能压力。激增的训练计算量模型参数规模可达万亿级别加之需要处理极长的用户行为序列导致训练计算量大幅增加对算力提出了极高要求。为应对这些挑战昇腾RecSDK提供了针对性的解决方案。在存储与通信方面其PyTorch生态的10TB分布式多级特征存储支持哈希映射、动态扩容、CPU Offload等功能旨在有效管理海量动态Embedding。在稠密计算优化方面它提供了关键算子如HSTU的优化实现支持变长序列、KV缓存、显存优化等特性以提升训练和推理效率。生成式推荐训练demo为了减少用户学习成本昇腾RecSDK通过开源recsys-examples修改接口运行到昇腾设备。 项目地址https://gitcode.com/Ascend/RecSDK/blob/develop_torch_benchmark/torch2.6.0_examples_benchmark/develop/gr_nv/README.md硬件准备硬件Atlas 800T A2 1台准备运行环境下载基础镜像https://www.hiascend.com/developer/ascendhub/detail/9faeb4847b3e419f81b78a4d0ed574b5运行容器命令参考#!/bin/bashcontainer_name$1image_name$2dockerrun\-it\--name${container_name}\-eASCEND_VISIBLE_DEVICES0-7\--shm-size300g\-v/etc/localtime:/etc/localtime:ro\-v/etc/ascend_install.info:/etc/ascend_install.info:ro\-v/usr/local/Ascend/driver:/usr/local/Ascend/driver:ro\${image_name}\/bin/bash依赖软件版本如下Python版本主要配套依赖Python3.11torch2.6.0torch_npu2.6.0fbgemm_gpu1.1.0cputorchrec1.1.0npuhybrid_torchrec1.1.0torchrec_embcache1.1.0运行安装依赖项下载https://gitcode.com/Ascend/RecSDK/blob/develop_torch_benchmark/torch2.6.0_examples_benchmark/develop/gr_nv/README.md安装 gin-configpip3installgin-config参考 https://gitcode.com/Ascend/MindSpeed/tree/core_r0.8.0 安装 MindSpeed 和 Megatron-LM模型库源码适配进入当前目录下载官方模型代码后并使用patch文件进行修改。gitclone https://github.com/NVIDIA/recsys-examples.gitcdrecsys-examplesgitcheckout v25.05cp-f../gr_nv2npu.patch ./gitapply gr_nv2npu.patch模型运行拷贝run.sh到hstu目录cp../run.sh ./examples/hstu/cd./examples/hstu参考https://github.com/NVIDIA/recsys-examples/blob/v25.05/examples/hstu/README.md#dataset-preprocessing 准备数据集(示例为ml-20m数据集)。执行命令bashrun.shRecSDK如何支撑生成式推荐训练昇腾RecSDK面向传统推荐、生成式推荐与多模态推荐等多样化场景基于高性能、高可靠的开发者工具链构建为生成式推荐训练提供全面支持。其整体架构如下大规模分布式稀疏Embedding表提供与TorchRec兼容的接口无缝对接开源生态并补充支持动态哈希、多级缓存、数据/模型并行策略、多级流水线、梯度累积以及特征准入与淘汰机制有效应对生成式推荐场景下海量特征ID的Embedding存储需求。关键算子的昇腾原生实现涵盖PyTorch生态中Fbgemm相关必要算子以及专为昇腾优化的HSTU融合算子、查表算子和通用计算融合算子充分发挥硬件算力。基于MindSpeed的分布式并行支持依托昇腾PyTorch生态的MindSpeed分布式并行框架实现HSTU侧的并行计算能力提升训练效率。依托MindStudio的DFX能力基于昇腾CANN生态的MindStudio工具链提供系统化的性能分析与精度调优手段保障模型训练过程的可靠性。支持自定义算子开发通过AscendC编程架构开放灵活的自定义算子扩展能力满足个性化算法需求。HSTU加速算子HSTUHierarchical Sequence Transduction Unit融合算子是生成式推荐模型的核心计算组件其针对昇腾硬件架构进行了深度优化。具体见链接实现如下功能基于分离架构的Cube/Vector计算并行依托昇腾AI Core的分离式计算架构昇腾社区HSTU融合算子实现了Cube矩阵计算单元与Vector向量计算单元的并行流水线执行。该设计允许在Cube单元执行块状矩阵运算的同时Vector单元可处理相应的Silu/Add/Muls等计算从而有效掩盖计算延迟显著提升硬件算力的整体利用效率。基于Block切分的内存复用与带宽优化采用动态Block切分策略将大尺寸张量计算分解为适宜硬件缓存的小数据块。此举不仅适配了片上缓存容量更重要的是实现了数据在局部缓存L2中的高效复用大幅减少了访问全局内存的频次与数据搬运量从而有效降低了内存访问延迟。并且在搬运一个Block的同时能够进行另外一个Block的计算从而提升计算带宽的利用率。贪心均匀分核策略实现负载均衡在面对不规则或动态变化的数据负载时如用户行为序列HSTU融合算子采用贪心均匀分核策略进行任务调度。该策略旨在将计算任务尽可能均匀地分配到多个计算核心上减少了因单个核心负载过重或过早空闲带来的等待时间确保了多核之间的负载均衡从而充分发挥分布式计算的优势。支持复合mask的跳过策略。为优化长序列处理中的无效计算该算子支持复合注意力Mask机制。通过识别序列中的无效或填充位置并依据特定规则如因果掩码、Context/Target Mask跳过生成复合Mask从而在计算过程中智能跳过对无效位置的计算。此策略在特定场景下如处理高度稀疏的序列数据可显著减少计算量例如减少超过50%的QK相关计算量。稀疏表框架基于Torchrec的API接口链接提供EmbeddingCollection与EmbeddingBagCollection建表功能并提供TrainPipelineSparseDist实现CPU计算、H2D/D2H、NPU计算、NPU通讯之间的并行。通过Swap引擎和EmbCache管理模块与上层训练交互其中Emb存储模块的核心目标是在CPU上高效、线程安全地管理大规模Embedding参数的存储、包括查找、插入、更新和删除等操作。自研高性能哈希表 FastHashMap并通过统一的抽象接口 EmbTable 对外屏蔽底层细节。为进一步提升性能和内存利用率引入了内存池EmbMemoryPool 和自定义分配器支持embedding的高效分配与回收。灰色CPU资源蓝色NPU计算资源绿色H2D和D2H拷贝资源为了优化CPU、通信与NPU之间的协同工作充分挖掘昇腾硬件的计算潜力RecSDK设计并实现了一套高效的流水并行策略。该策略的核心思想是在一个训练步骤Step中预先获取多个批次Batch的数据并将不同批次数据处理的多个任务阶段进行重叠执行从而形成流水线式的并行处理显著提升整体吞吐率。如图所示RecSDK将训练任务精细地划分为三种不同类型的并行执行流Device计算流这是任务执行的主干流水线主要负责在NPU上完成核心计算任务包括但不限于Embed在设备内存中的换入换出操作、模型的前向传播计算以及反向传播过程中的梯度计算和参数优化更新。CPU计算流此流水线在CPU的线程上运行主要负责Host侧的预处理与逻辑控制任务。典型操作包括Embedding向量的初始化、基于索引的查找、参数的更新维护以及为Device计算流中的换入换出操作提供位置查询信息等。拷贝流该流专职负责Host侧CPU内存与Device侧NPU内存之间的数据搬运和交互确保计算所需的数据能及时送达NPU并将计算结果高效回传。通过预先获取多个Batch的数据RecSDK的流水并行策略使得上述三个流程能够并行不悖地工作。例如当第N个Batch的数据正在Device计算流上进行模型训练时第N1个Batch的数据在拷贝流上进行传输而第N2个Batch的数据则在CPU计算流上进行预处理。这种任务级并行极大地减少了各个硬件组件CPU、NPU和数据传输链路通信的闲置等待时间从而实现了训练性能的显著提升。结论与未来规划为了支撑昇腾用户在生成式推荐领域的探索与实践RecSDK提供了与主流开源模型高度兼容的API接口设计旨在帮助开发者无缝迁移现有工作流并快速构建和部署生成式推荐模型。仓库见链接文档见链接。针对生成式推荐任务面临的高存储消耗与大规模计算挑战RecSDK重点引入了多项关键特性在稀疏侧提供了高效的大规模分布式稀疏表在稠密侧集成了经过深度优化的HSTU融合算子。这些特性共同助力生成式推荐模型在昇腾硬件上获得极致的性能表现。展望未来我们将持续完善RecSDK的功能体系进一步弥合与开源生态的功能差异拓展对更多主流生成式推荐模型的适配范围并通过底层技术的持续创新为用户提供更极致的性能优化。