1. 项目概述当推荐系统撞上内存墙我们如何破局在当今的互联网服务中个性化推荐系统是提升用户体验和商业价值的关键引擎。无论是电商平台的“猜你喜欢”还是内容平台的“信息流”其背后往往都运行着复杂的深度学习推荐模型。然而随着模型参数规模的爆炸式增长尤其是嵌入层动辄数十甚至数百GB的参数量传统的计算架构正面临严峻的“内存墙”挑战。简单来说模型太大数据搬不动了。CPU内存带宽有限成为拖慢整个推理流程的瓶颈而GPU虽然算力强大但其昂贵且容量有限的内存又难以容纳庞大的嵌入表。这种矛盾在追求低延迟、高吞吐的在线推理场景中尤为突出。GPAQuant正是为了解决这一核心矛盾而诞生的。它并非一个天马行空的理论构想而是一套基于商业现货COTS处理内存PIM硬件深度融合模型量化技术的实战型加速框架。其核心思想直击痛点既然数据搬运慢那就让计算去找数据既然内存装不下那就把数据“压缩”后再装。通过将量化后的嵌入表分布式存储在UPMEM DPU的高带宽内存中并将耗时的嵌入查找和归约操作卸载到DPU上执行GPAQuant在嵌入层处理上实现了最高4.6倍的加速端到端推理实现了最高3.4倍的加速且精度损失几乎可以忽略不计。这篇文章我将从一个系统工程师的视角深入拆解GPAQuant的设计精髓、实现细节以及我们在实际探索中踩过的坑和收获的经验希望能为面临类似性能瓶颈的同行提供一条可复现、可落地的优化路径。2. 核心挑战与设计思路拆解在深入技术细节之前我们必须先理解要解决什么问题以及为什么选择PIM量化这条技术路线。这决定了后续所有优化手段的方向和边界。2.1 传统架构的瓶颈到底在哪主流的DLRM推理部署通常采用CPU-GPU混合架构。嵌入表因其巨大的体积常为数十GB而被存放在CPU主存中而计算密集的全连接层则在GPU上执行。这个架构的瓶颈非常清晰内存带宽瓶颈CPU的内存带宽例如DDR4-2666四通道约85GB/s相对于嵌入查找的随机、稀疏访问模式而言严重不足。嵌入层操作本质上是一系列不规则的、高并发的内存读取极度依赖高带宽。我们的性能剖析显示在典型工作负载下嵌入层操作可占据总推理时间的90%以上CPU成了整个流程中最慢的一环。数据搬运开销每次推理都需要在CPU和GPU之间搬运嵌入向量引入了不可忽视的PCIe通信延迟。尽管有研究尝试在GPU上缓存热点数据但缓存命中率受数据分布影响大且无法从根本上解决容量问题。容量与成本的矛盾试图用GPU HBM来存储全部嵌入表成本极高且容量增长追不上模型规模的膨胀速度。2.2 为什么是PIM 量化面对上述瓶颈业界探索过多种方案如近内存处理Near-Memory Processing、使用SSD作为扩展内存等。GPAQuant选择基于UPMEM PIM进行软硬件协同设计主要基于以下几点考量突破性的内存带宽单个UPMEM DIMM模块能提供高达7.2 TB/s的聚合内存带宽这比传统CPU内存带宽高出两个数量级。这为嵌入层这种内存密集型操作提供了根本性的解决方案。真正的“内存中处理”UPMEM的DPU核心直接与专属的64MB MRAM内存库相连指令和数据存取都在芯片内完成彻底消除了数据在内存控制器和计算核心之间搬运的开销。商业可行性与能效比与基于ReRAM、FPGA等定制化PIM方案相比UPMEM是已经商业化的产品具有明确的编程模型和工具链落地可行性更高。同时其能效比TCO相比GPU有数量级优势这对于大规模部署至关重要。量化的必要性PIM解决了带宽问题但单DPU 64MB的内存容量对于巨型嵌入表仍是杯水车薪。因此必须引入模型量化将FP32的嵌入向量压缩为INT8甚至INT4才能将一张大表切分后放入多个DPU的MRAM中。然而量化会引入计算开销反量化和精度损失这要求我们必须设计PIM感知的量化策略。2.3 GPAQuant的整体设计哲学GPAQuant的设计遵循一个核心原则在硬件的约束下寻求最优的软件映射。它不是简单地把算法扔到PIM上跑而是深刻理解UPMEM DPU的架构特性如无直接DPU间通信、浮点运算慢、MRAM访问对齐要求等和工作负载特征嵌入访问的幂律分布进行针对性的协同优化。整个框架分为预处理和推理前向传播两个阶段预处理阶段离线在CPU上执行一次。核心任务有两项一是如何将庞大的、量化后的嵌入表智能地“切片”并分配到各个DPU上分区优化二是如何量化能既减少存储又最小化在DPU上反量化的开销量化优化。推理阶段在线系统接收稀疏输入特征根据分区结果将查找请求分发到对应的DPU。每个DPU在其本地MRAM中执行整数格式的嵌入查找和聚合尽可能推迟浮点反量化操作最后将部分和返回CPU进行最终合并。这个流程听起来直观但魔鬼藏在细节里。接下来我们就深入这两个核心优化部分。3. 分区优化在DPU的方寸之间排兵布阵将一张可能上百GB的嵌入表分割成许多个小于64MB的“瓦片”Tile并分配到数百个DPU上这就像一场高难度的拼图游戏。目标不仅是“装得下”更要“跑得快”、“负载匀”。3.1 架构感知分区理解硬件的“脾气”第一层优化是架构感知分区。这里我们需要像硬件工程师一样思考。UPMEM DPU有几个关键约束容量限制每个Tile必须 ≤ 64MB。对齐访问MRAM读取必须以8字节对齐单次最大读取2048字节。未对齐或过小的读取会造成带宽浪费。无直接通信DPU之间不能直接交换数据任何跨DPU的数据交换都必须通过CPU内存中转开销巨大。因此分区的目标是在满足容量约束的前提下最小化总耗时总时间 CPU-DPU通信时间 DPU查找时间 DPU-CPU通信时间。一个关键的发现来自于对MRAM读延迟的微观测试如图4所示。数据显示当读取数据大小从8字节增加到32字节时延迟增长非常缓慢但超过32字节后延迟急剧上升。这给了我们一个强烈的设计暗示应尽量让每次嵌入向量读取的数据量落在8-32字节的“甜蜜区”内。假设我们使用4比特INT4量化每个标量占0.5字节。一个嵌入向量维度为Nc加上存储缩放因子scale和偏置bias的4字节FP16总大小为Nc * 0.5B 4B。对齐到8字节后我们期望这个值 ≤ 32B。由此可以推导出Nc的最佳取值通常是2、4、8、16、32这些2的幂数。通过穷举搜索我们可以为给定的总行数R、列数C和DPU数量Ndpu找到使总时间最小的Nr每Tile行数和Nc组合。实操心得在实际编码中这个搜索过程可以很快。我们通常预先计算好不同(R*C, Ndpu)组合下的最优(Nr, Nc)配置表在系统初始化时直接查表避免在线计算开销。3.2 工作负载感知分区应对“热点”的智慧均匀分区假设每个嵌入向量的访问频率是相同的但这与现实严重不符。推荐系统中的物品访问遵循典型的幂律分布即少部分热门物品Hot Items被绝大多数请求访问。如果简单按行号范围均匀分区那么存储了热门物品的DPU就会成为性能瓶颈其他DPU则闲置无法发挥PIM的并行优势。因此我们引入了工作负载感知分区。其核心思想是在分配Tile到DPU时不仅要考虑存储空间还要考虑各Tile的预期访问频率追求各DPU总访问负载的均衡。算法1描述了这一过程。输入是物品的访问频率历史obj_freq和由现有缓存技术如GRACE生成的缓存结果cache_res。GRACE会分析数据找出经常共现的物品组合例如{用户A 电影B 类别C}并预计算它们的部分和进行缓存从而将多次查找合并为一次。我们的分区算法是一个贪心策略将物品按访问频率降序排序。对于缓存列表中的物品组合我们将其整体分配到一个当前总负载最低且有足够缓存空间的DPU上。这相当于把“热点群”打包处理充分利用缓存收益。对于不在缓存列表中的物品将其分配到当前总负载最低且有足够嵌入表空间的DPU上。这样通过将高频访问的Item和其缓存友好组合智能地分散到不同DPU我们实现了工作负载的均衡。如图8所示当请求查找索引1,4,5时索引4和5的预计算和45可以从DPU0的缓存中一次性读出而索引1则从DPU1的嵌入表区读取。两个DPU并行工作最后CPU只需对两个部分和做一次加法。避坑指南缓存空间和嵌入表空间共享64MB的MRAM。这里需要一个权衡缓存空间越大缓存命中带来的加速收益越高但可用于存储嵌入向量的空间就越小可能导致需要更多DPU。我们的实验表明在Goodreads数据集上当缓存空间占所需缓存列表大小的100%时相比无缓存方案嵌入查找时间可减少26%。通常我们会将缓存容量设置为一个可配置参数根据数据集的热点分布程度进行调整。4. 量化优化在精度与效率的钢丝上行走量化是压缩模型、适配PIM内存容量的必要手段但在PIM上量化面临着特殊挑战DPU的浮点运算性能远弱于整数运算。一次标准的反量化操作dequant(x_int) x_int * scale bias包含一次浮点乘法和一次浮点加法。如果对每个查找回来的嵌入向量都立即反量化那么在DPU上的计算开销将变得不可接受。4.1 分块量化与瓦片内重排推迟反量化的艺术为了解决这个问题GPAQuant的核心策略是推迟反量化。我们不是对每个向量单独反量化后再相加而是先对多个量化后的整数向量求和再对总和进行一次反量化dequant(x_int1 x_int2)。这能将浮点操作次数从O(n)降低到O(1)。但这要求被加和的向量必须共享相同的量化参数scale和bias。因此我们不能采用传统的逐行量化每行有自己的参数而必须采用分块量化即一个Tile或Tile内的一个组内的所有行共享一套量化参数。然而分块量化会扩大量化误差。因为块的scale必须覆盖组内所有值的范围范围越大根据公式|e(x)| ≤ scale/2误差上限就越大。为此我们提出了瓦片内重排算法Algorithm 2。其目标非常直接通过重新排列一个Tile内行和列的顺序使得被分到同一组内的数值其范围尽可能小。算法分为两步行排序对于Tile的每一列切片计算每行元素的和并按和降序排列。这样数值大小接近的行会被聚集在一起。列重排这是算法的精髓。对于Tile的每一行切片我们计算每列的最小值C_min、最大值C_max和范围disc。然后我们尝试将列分组使得同一组内的所有列其值域范围[C_min[i], C_max[i]]都落在某个共同的区间[cur_min, cur_max]内。我们从范围最宽的列开始逐步放宽cur_min和cur_max的约束将尽可能多的列纳入同一个组。被分到同一组的列将共享量化参数而由于它们的值域高度重叠所需的scale就会很小从而显著降低量化误差。如图9的例子所示经过行列重排后所有Tile的scale值之和减少了三分之一这意味着整体量化误差得到了有效控制。4.2 自适应分组量化当一刀切不再适用瓦片重排解决了组内问题但Tile本身的大小是由分区策略决定的可能差异很大。对一个巨大的Tile使用同一套量化参数误差仍然可能很大。因此我们需要在Tile内部进行更细粒度的划分即分组量化。但组的大小如何确定组太小如每行一组反量化开销大组太大量化误差大。GPAQuant的创新在于使用一个轻量级梯度提升决策树GBDT模型来动态决定每个Tile的最佳分组大小。在预处理阶段我们从每个Tile中提取三个元特征如表III所示F0: 嵌入表总大小。F1: 该Tile的scale与一个基准最小scale的比值。F2: 整个嵌入表的全局scale与基准最小scale的比值。这些特征反映了数据的分布特性。我们使用部分数据集如Kaggle, Criteo进行训练通过网格搜索为不同的Tile找到在精度损失最小情况下的最优分组大小如128, 256, 512行。训练好的GBDT模型在预处理时对每个Tile进行快速推理决定其分组策略。由于这个推理只在预处理阶段执行一次它对线上推理的延迟是零开销的。经验之谈GBDT模型在这里的成功应用关键在于特征工程。我们最初尝试了更多统计特征如均值、方差、峰度但发现反而增加了模型复杂度且提升有限。最终选择的这三个抽象特征既能表征数据分布又保证了模型的轻量和泛化能力。训练一次模型仅需几分钟CPU时间但可以泛化到不同的数据集和模型结构上。5. 系统实现与性能调优实录将理论设计转化为实际可运行的系统是工程上最具挑战性的一环。这里分享我们在实现GPAQuant过程中的一些关键细节和调优经验。5.1 DPU内核编程与优化UPMEM DPU使用一个定制化的ISA编程模型类似于多线程RISC核心。我们需要为嵌入查找和整数聚合编写高效的内核代码。内存访问模式优化DPU的MRAM带宽虽高但必须通过WRAM便签式内存进行访问。我们的内核设计采用双缓冲技术当一组tasklet在处理WRAM中当前数据块时另一组tasklet通过DMA引擎预取下一个数据块到WRAM。这有效地隐藏了内存访问延迟。整数聚合与延迟反量化内核的核心循环是整数向量的加法。我们确保所有参与一次归约的向量索引都来自同一个量化组。在代码中这体现为在查找前先根据索引判断其所属的组ID。只有属于同一组的向量才进行整数累加否则需要立即反量化并转换为浮点数进行累加。这部分逻辑需要精细设计以避免分支预测失败带来的性能损失。Tasklet并行度每个DPU支持最多24个tasklet硬件线程。我们通过实验发现对于嵌入查找这种内存密集型任务将tasklet数量设置为14左右能达到带宽利用率和线程管理开销的最佳平衡点。过多的tasklet反而会因资源竞争导致性能下降。5.2 主机端调度与通信重叠主机CPU负责协调整个推理过程。其性能关键在于最大化DPU的利用率和重叠通信与计算。请求批处理与分发CPU将一批推理请求的稀疏索引收集起来根据预处理阶段生成的分区映射表将索引拆分并打包成不同的消息通过DMA批量发送到对应的DPU。这里消息的大小需要对齐到DPU的访问粒度我们通常设置为32字节的倍数。异步通信与计算我们使用异步DMA操作。CPU在发出所有DPU的数据传输命令后并不等待而是可以继续处理之前批次的DPU返回结果或者准备下一批数据。UPMEM SDK支持这种异步编程模式充分利用了PCIe总线的双向带宽。结果归约各个DPU返回的是部分和可能是整数也可能是浮点数。CPU需要将这些部分和按原始请求的顺序进行最终归约。我们使用OpenMP并行化这一归约过程特别是当批量大小较大时并行归约能带来显著收益。5.3 精度与速度的权衡实践在真实业务中我们永远在平衡精度和速度。GPAQuant的量化策略提供了多个可调节的旋钮量化比特数论文中主要使用INT4我们在内部也测试了INT8。INT8的精度损失更小有时甚至无损但存储和带宽压力是INT4的两倍。选择取决于业务对精度的容忍度。一个实用的方法是在离线评估中绘制一条“精度-比特数”曲线找到那个精度陡降的拐点拐点之前的比特数就是性价比最高的选择。分组大小GBDT模型预测的分组大小是一个基础。在实际部署中我们可以设置一个分组大小的上限和下限。例如即使模型预测某个Tile适合512行的大组但如果该Tile包含某个极其重要的特征域如“用户ID”我们可能会强制将其分组大小限制在256行以换取更精确的表示。分层量化对于超大规模的嵌入表我们实践过一种“分层量化”策略对访问频率最高的前1%的嵌入行使用INT8甚至FP16存储其余部分使用INT4。这需要在分区时将这些“特权行”集中分配到少数DPU并修改内核逻辑进行特殊处理。虽然增加了系统复杂性但在一些对头部物品推荐精度极其敏感的场景下效果显著。6. 效果评估与生产环境考量我们依据论文中的实验设置在六个真实数据集上进行了全面的评估对比了DLRM-CPU、HotlineCPU-GPU混合方案、PIM-Uniform仅均匀分区和GPAQuant。6.1 性能数据解读嵌入层加速如图10a所示GPAQuant在嵌入层操作上取得了压倒性的优势。相对于DLRM-CPU加速比最高达4.6倍。更重要的是在访问热度很高的数据集如GoodReads上GPAQuant相比不考虑负载均衡的PIM-Uniform仍有高达2.3倍的提升这充分证明了工作负载感知分区的价值。端到端加速如图10b所示端到端加速比达到3.4倍。这意味着GPAQuant不仅解决了嵌入层瓶颈其高效的设计也使得其他层如MLP的执行能更快开始减少了整体的流水线空泡。能效比功耗方面两个UPMEM模块总功耗约46W而对比的CPU TDP为85W。结合大幅缩短的运行时间GPAQuant能实现46%-88%的能耗节省。对于数据中心规模的应用这意味著巨大的电费成本和碳排放的降低。6.2 生产环境部署经验将研究原型推向生产环境会面临更多工程挑战冷启动与模型更新预处理阶段分区、量化、重排是需要时间的。对于超大规模模型这个过程可能需要数小时。在生产中我们采用“蓝绿部署”策略在新版本模型预处理的同时旧版本模型继续服务预处理完成后瞬间切换流量。对于实时性要求极高的特征更新我们设计了增量更新机制只对发生变化的部分嵌入行进行重计算和重新分发。容错与弹性UPMEM DIMM或单个DPU可能出现故障。我们的系统在主机端维护了DPU的健康状态。一旦检测到某个DPU失效可以动态地将分配给它的Tile重新分配到其他健康的DPU上当然这会牺牲一些负载均衡性。同时我们保留了在CPU上运行嵌入层的后备代码路径在PIM硬件完全不可用时可以降级运行。监控与调试我们开发了一套详细的性能指标收集系统可以监控每个DPU的负载、缓存命中率、量化/反量化操作次数、通信流量等。这些指标对于定位性能热点、调整分区和量化参数至关重要。特别是当业务数据分布发生漂移时这些监控能预警是否需要重新进行预处理。6.3 常见问题与排查技巧在实际运维中我们遇到并总结了一些典型问题性能未达预期检查点首先用nvidia-smi和UPMEM工具查看GPU和DPU的利用率。如果DPU利用率低很可能是工作负载不均衡。使用我们的监控工具查看各DPU的请求处理时长分布。排查方法导出分区映射表分析热门物品是否集中在少数DPU。如果是检查预处理时输入的访问频率日志是否具有代表性或尝试调整缓存分配策略。精度下降超出预期检查点首先在离线环境下用全精度模型和GPAQuant量化模型在同一份测试集上对比AUC/LogLoss。排查方法如果精度下降明显重点检查量化环节。可以逐层输出量化前后的嵌入向量计算误差。通常问题出在少数值域范围特别大的特征域上。可以为这些特征域单独设置更小的分组大小或更高的量化比特位。DPU内核执行错误检查点UPMEM SDK提供了调试工具可以单步执行DPU内核。错误通常源于内存访问越界或数据类型转换溢出。排查方法在内核中大量使用printf输出到主机进行调试。特别注意整数累加时的溢出问题INT4的累加和可能超出INT32的范围需要适时提升到INT64。7. 未来展望与扩展思考GPAQuant的成功验证了基于商业PIM进行软硬件协同设计是一条行之有效的路径。放眼未来我们认为还有几个值得深入的方向支持训练目前GPAQuant专注于推理加速。如何将PIM的高带宽特性用于嵌入层的分布式训练特别是处理稀疏梯度更新是一个更大的挑战。这需要设计新的通信原语和一致性协议。异构内存层次未来的系统可能是PIM、HBM、DRAM、CXL扩展内存共存的异构体。如何智能地将嵌入表的不同部分如根据访问频率、更新频率放置在最合适的存储层级上并实现无缝的数据迁移是一个系统级的优化问题。算法-硬件协同设计GPAQuant是在现有DLRM模型和PIM硬件之间的适配。更激进的想法是从模型结构设计之初就考虑PIM的约束。例如设计更适合分块量化的嵌入表示方法或者探索基于PIM的新型推荐模型架构。回过头看GPAQuant项目的最大收获不仅仅是性能数字的提升更是一种思维模式的转变从“让硬件适应算法”到“让算法与硬件共舞”。在算力瓶颈日益凸显的今天这种深度的软硬件协同优化将是解锁下一代AI系统性能的关键。