单双层架构解析

Meta GR 等系列研究揭示，推荐系统的性能同样遵循 Scaling Law，即通过扩展模型规模（尤其是CTR模型中的稀疏与稠密参数）能够带来效果的持续提升。这一规律正驱动着京东等互联网企业积极布局大规模推荐系统的研发。如下表所示，为应对十亿至万亿级别的稀疏特征，字节跳动、谷歌与京东等公司的Embedding表规模已激增至数千GB乃至PB级别。如此庞大的数据体量对存储与访问架构构成了严峻挑战，其中巨大的内存消耗与特征量的指数级增长，已成为制约模型规模扩展的主要瓶颈。

为应对上述挑战，业界演化出两种核心架构：

单层架构：将Embedding表全量数据存放于设备侧（如NPU），依托设备间高速通信实现协同计算。该架构以访问效率为核心优势，适用于数据规模可控、可完全载入设备内存的场景，无需考虑数据的换入换出。如下图右半部分所示，在单层架构中，整个Embedding表被加载至设备侧，由NPU统一负责计算执行：首先对batch数据进行预处理，创建多表Embedding collection并进行分片；随后通过lookup和parse lookup等操作获取对应的嵌入向量；最终完成稠密模型的前向计算、稀疏与稠密参数的反向传播以及梯度计算等流程。值得注意的是，该计算流程在单层与双层架构中基本保持一致。

双层架构：采用 “Host 侧 + Device 侧” 分层存储，Host 侧通过EmbCache Manager 管理全量 embedding 表，对输入的 KeyedJaggedTensor 执行哈希、分桶等预处理，将数据暂存于 DDR 中，并负责参数初始化与key值的查询。Device 侧仅加载当前计算所需数据。其核心是突破 Device 内存限制，解决千亿级特征带来的存储压力，支持动态扩容。数据换入换出、动态扩容这一些列操作是双层架构为解决 Device 内存不足而设计的。

如流程图右侧所示为单层架构流程，整个Embedding表置于设备内存。如图的整体框架所示为双层架构流程，通过在Host侧引入主机内存进行分层存储，以突破设备内存的容量限制：

因此，单双层架构并非替代关系，而是互补关系。正确的架构选型策略是：在Embedding表规模可控的场景下采用单层架构以追求极致的性能；而当模型规模扩展到单层内存无法容纳时，则采用双层架构，在可接受的性能损耗下实现模型的规模持续扩展。

1. 动态扩容

在实际训练推荐系统模型时，嵌入表（Embedding Table）的整体规模通常非常庞大，但其中仅有少部分特征数据会被频繁访问，存在显著的长尾分布现象。若某些特征出现的频率过低，不仅对模型训练效果的提升无益，还会造成存储资源的浪费，甚至因引入噪声而导致模型过拟合。在传统模式下，嵌入表的尺寸需在初始化时固定，无法在训练过程中调整，容易导致显存资源分配不足或浪费。而在推荐场景下，多个稀疏表的实际规模往往难以预先准确估计。为此，Rec SDK提供了动态扩容功能。用户可通过设置参数 use_dynamic_expansion = True 启用DDR内存侧的动态扩容模式，以更好地适配实际需求。

具体来说，动态扩容支持两种资源层面的扩展方式：

● 片上内存（On-Chip Memory）动态扩容：指显存容量随训练进程动态增长，适用于嵌入表完全载入显存的场景；

● DDR动态扩容：指对主机内存占用随训练数据增长，而显存占用不变，适用于超大规模稀疏表的多级存储架构。

值得注意的是，片上内存显存侧动态扩容模式下，不支持特征淘汰。因此在实际配置时，需根据是否开启特征淘汰功能合理选择扩容策略。

动态扩容的触发流程：

1. 触发条件：当用户调用插入接口（如 find_or_insert 或 insert_or_assign）尝试向嵌入表中添加新的特征键（Key）时，系统会实时计算当前的负载因子。
2. 触发扩容：如果 负载因子 ≥当前存储的元素数量 / 总容量，则意味着当前哈希表已趋于饱和，发生哈希冲突的概率将显著增加，进而导致查询和插入性能下降。此时，系统会自动触发扩容（double capacity）流程

find_or_insert

接口说明:find_or_insert 为用户提供查询并插入操作，如果在HBM侧表中能够查到指定的Key值则返回其对应的嵌入向量和优化器状态，如果查不到则插入新的keys。如下图所示为该接口的关键流程：

● step1: hostcpu 根据当前负载判断是否需要进行扩容。

● step2: 如果需要扩容，则host分配DDR的bucket内存。然后通过allocate_bucket流程将bucket内存进行分配。

● step3: 如果需要扩容，bucket内存分配之后，需要进行将old key进行rehash操作。

● step4: 调用find_or_insert 对 new key 进行查找并插入操作，并初始化其嵌入权重。

insert_or_assign

接口说明:为用户提供key（键值对）插入并赋值到嵌入表的接口，如果key已经存在则赋值(更新)，如果key不存在则插入，返回插入失败的key在输入的索引位置。该接口的关键流程如下:

● step1:用户调用insert_or_assign接口指示需要在table_id对应的hashtable中需要插入的keys和values以及scores。

● step2:调用算子insert_or_assgin进行插入操作，如果key在hashtable中可以找到则进行更新操作，如果key在hashtable中找不到则进行插入操作。

● step3:插入的过程中，如果bucket非空则线性探测找到空位插入，如果bucket满了则比较分数，淘汰掉较小的分数对应的key，如果待插入得key的score比当前bucket中所有key的score都小，则插入失败。

2. 准入与淘汰

在大规模的稀疏场景下的模型训练，部分特征频次较低，无法为模型的训练提供有效的信息，同时会造成内存的浪费。也存在部分长时间不进行更新的特征，由于其时效性低，干扰训练结果。因此我们针对业务场景提出合适的准入淘汰机制，对动态场景下的词表进行精细化管理。目前我们的准入淘汰功能在CPU-Ps中进行控制，前向阶段进行准入计数和准入判断，反向阶段进行淘汰计数，在save阶段或固定step阶段执行淘汰。

2.1 准入策略

在推荐系统和机器学习中，特征准入机制是管理稀疏特征（如用户ID或物品ID）的关键技术，用于控制哪些特征键（key）被纳入模型训练或推理，以优化内存使用、计算效率和模型性能。默认采用NoneFilters，即所有key均准入。如下所示为目前支持三种准入策略：

基于访问次数： key出现次数>=阈值，则该key被准入。这类似于低频过滤，用于淘汰不常见或噪声特征，减少存储开销。

基于概率： 生成随机数 < 阈值（[0,1]之间）则key被准入。这样避免模型偏差于高频特征，提升多样性。

基于show-click： 结合用户行为数据（如展示次数show_cnt和点击次数click_cnt），计算ShowClick score 。

ShowClick score = alpha ** show_cnt + beta **  click_cnt（alpha，beta为自定义参数）。

如果得分超过阈值，则key被准入。这种机制适用于CTR预估或广告推荐场景，优先保留有高交互概率的特征，从而提升模型相关性。

这些机制可单独或组合使用（如先基于访问次数过滤，再通过概率机制探索），并根据业务需求调整阈值。整体上，准入策略旨在平衡特征覆盖率与计算效率，尤其在处理高基数稀疏数据时至关重要。

2.2 淘汰策略

在推荐系统的稀疏特征表管理中，淘汰策略是优化内存使用和模型性能的核心机制之一。当系统资源受限或需要维护特征表的高效性时，合理的淘汰策略能够自动移除低价值特征，确保高频重要特征得到保留。默认采用NoneShrink，即所有key均不淘汰。如下所示为目前支持的五种淘汰策略：

基于固定step： 训练过程中，key在连续多个训练步长（step）内未参与训练（即未被访问或更新），则触发淘汰。

基于访问次数： 每个key维护一个version变量，表明特征的新鲜程度。 若version >= threshold，且代表这个key很久没有用 过，则予以淘汰。

基于L2范数： 计算所有key的L2范数（即各维度平方和的平方根）L2_score，若L2_score < threshold（反映特征权重整体较小），则认为该特征对模型贡献微弱，触发淘汰。L2范数计算公式如下图所示：

基于时间和频次： 时间time：key被尝试淘汰的次数；频次freq：key参与训练的总次数；

在出现频次最低的k个key中，如果这个key属于尝试淘汰频率最高的k个key，则淘汰这个key；

基于show-click： 针对广告或推荐场景，key的评分由展示次数（show_cnt）和点击次数（click_cnt）加权计算：

ShowClick score = alpha * show_cnt + beta *  click_cnt（alpha，beta为自定义参数）。

在得分最低的k个key中，进一步筛选点击率（click-through rate）最低的key进行淘汰。

3. 优化器

在推荐系统的模型结构中，参数分为稀疏部分（sparse） 和稠密部分（dense） ，二者特性差异显著，因此需要配置不同的优化器。

其中，稀疏部分涉及表示词汇表的大规模 ID（如用户 ID、物品 ID），这类参数对内存需求极高。以 “10B 规模词汇表、512 维嵌入表、搭配 Adam 优化器” 为例，仅用 FP32 精度存储嵌入向量和优化器状态，就需要 60TB 内存，内存压力极大。为缓解这一问题，方案采用row-wise AdamW 优化器（参考 Gupta 等人 2014 年、Khudia 等人 2021 年的研究），并将优化器状态存储在 DRAM（动态随机存取存储器）中，此举可将每个浮点数（嵌入向量）的 HBM（高带宽存储器）使用量从 12 字节降至 2 字节。而模型的稠密部分参数量相对较小，无需特殊优化，直接使用 PyTorch 原生的torch.optim优化器即可。

从训练流程来看（如示意图所示），针对大规模稀疏表，每次训练时会先根据索引（index）从稀疏表中筛选出需要计算的若干行数据，再基于这些数据完成梯度计算与梯度更新操作。

优化器反向更新

如下图，清晰展示了一个针对稀疏嵌入参数的梯度处理与优化器更新的得流程。通过先对梯度按参数行进行排序和聚合，再启动优化器，实现大规模嵌入表的参数更新。

在当前 Unirec 框架中，嵌入层参数的更新是在反向传播过程中直接完成的。为实现优化器与稀疏表更新操作的整合以提升性能，并支持全局去重场景，框架采用了backward_codegen_adam_unweighted_exact这一融合算子 —— 它能结合 Adam 优化器与稀疏表的反向更新过程，根据不同优化器的特性计算并更新相关嵌入参数，同时负责维护优化器的状态。其类图结构可围绕这一融合算子与优化器、稀疏表的关联关系展开（如类图所示）。

4. 一表多查或多表单查

在嵌入表的 lookup（查表）流程中，输入参数已从原本的 feature ids 数据转换为 hash index（哈希索引）数据，再将这些索引传入算子执行查表操作，最终返回 embedding（嵌入向量）结果。

实际场景中，面对多个 slot（特征槽）的高效查询需求，主要有两种查表方式（均通过SplitTableBatchedEmbeddingBagsCodegen实现）：

● 多表单查：指对多个特征表进行一次查询。通过EmbeddingCollection在建表时合并维度相同的表实现，但查表时需额外信息（如下标所属表的标识、表的权重偏移等）来区分不同表的查询范围。

● 一表多查：指对同一张表进行多次查询，特点是 “前向过程多次查询，反向过程一次更新”。查询在 HBM（高带宽存储器）侧进行，若 HBM 中未找到所需数据，则通过换入机制从其他存储位置加载。

（具体流程可参考图示的嵌入表 lookup 过程）

下图清晰的展示了SplitTableBatchedEmbeddingBagsCodegen算子的查表流程。Hashembeddingbagcollection 作为基类用于创建模型，跑通查表更新流程。

如下图为一表多查的情况，中间维度为6的一个embedding表承担了多特征查询的任务，左侧的各个特征 ids，都能在这张中间的 embedding 表里，精准找到与之对应的向量表示，最终呈现出右侧的查表结果，实现了用单张 embedding 表高效完成多特征的嵌入查询。

如下图所示为多表单查的情况，将多张小表合并成一个大表，可以减少创建的embedding table个数。把下图右侧的 embedding table 1、embedding table 2、embedding table 3 这些小表合并成一个大表后，左侧一次查询的多个特征 ids，就能通过 “合表，查一次”，在这张大表中统一查询。

5. 模型save和load模块

由于Embedding嵌入层和稠密部分的数据格式不同，两者的 save（保存）和 load（加载）方式存在差异。其中，dense 部分直接采用 PyTorch 原生的 save 和 load 方式，而嵌入层的保存与加载则通过专门的模块处理，整体流程如下：

同时由于单双层的设计，导致单双层情况下embedding等数据的存储位置存在明显差值，导致save、load模块在单双层上的实现上也存在较大的差异。对于分片嵌入模块，目前双层使用EmbCacheShardedEmbeddingCollection类来获取参数和保存embedding，单层使用HybirdEmbeddingCollection类来获取参数和保存embedding。

目前框架的NPU支持保存、加载稀疏表，指定保存固定step的稀疏表，多个特征表独立管理，实现嵌入层的增量式的保存和加载，适用于保定期保存模型中间状态或恢复训练场景。

增量模型保存与加载

此功能是为了支撑流式训练开发，推荐系统在服务的过程中，会不断产生可用于训练CTR模型的日志数据，这些数据也是不断输入给模型。模型每接收一部分数据，就会利用该部分数据训练模型，同时按一定的频率（时间）保存全量或增量模型。

在这个背景下，将稀疏的参数仅保存增量模型，会极大的降低频繁保存模型带来的额外开销。同时，在训练进程结束后能尽量通过一个最近的全量的模型配合一系列的增量检查点（模型）恢复最近一次训练完成的模型参数，减少重复计算。

1. save模块

如下图右侧所示为单层save模块流程图，左侧所示为双层save模块流程图。

a) 通过不同分层模块获取设备端的优化器状态、动量信息和Embedding参数；

b) 双层架构数据转移，使用embedding_to_host将设备数据迁移至Host，嵌入内容会通过EmbcacheManager模块加载更新到FastHashMap。此时，Host侧保存的是embedding的全量信息，其中每个表都有自己对应的动量信息；

c) 断点续存支持，在Host侧对全量信息进行分桶保存，便于训练中断后能快速恢复。

d) 准入淘汰机制：双层情况下，同时保存准入淘汰的相关参数信息Shrinked keys，例如count，timestamps等信息。

如下图所示为保存的格式，其中：

global_step{step}：表示dense部分保存的文件目录

globalstep{step}{table_name}：表示sparse部分保存的文件目录

2. load模块

如下图右侧所示为单层load模块流程图，左侧所示为双层load模块流程图。

a) 参数加载，读取存储路径中.pt的内容hyper_param格式的优化器状态信息、学习率等信息；

b) 单层架构处理，读取存储的.data模块，加载每个表的动量信息，更新对应的embedding table；

c) 双层架构处理，调用EmbcacheManager的C++接口，在Host侧，对全量信息进行加载；

d) 数据还原，在Host侧加载每个表的动量信息，对分桶保存的嵌入内容进行还原映射，将嵌入内容加载到FastHashMap。如需准入淘汰机制的参与，同步在Host侧加载准入淘汰涉及的参数信息。