Fork me on GitHub

推荐广告模型的降本提效:压缩策略

背景

从全局看,深度学习模型规模在过去数年持续的指数膨胀。在模型效果提升的同时,为训练和推理的性能和成本都带来了严峻的挑战。作为应对,出现了大量的模型压缩策略,比如Zero系列的训练时策略。推理时量化、剪枝策略。以及训练推理协同策略,比如蒸馏等等。

广告、推荐的深度学习模型的体积也非常庞大,从GB到TB不等。然而这类模型比较特殊在于通常99%以上的参数都在embedding层。广告、推荐模型在许多互联网公司占用了绝大部分的AI算力开销,超过了于CV、NLP等其他模型。主要有几方面原因:

  • 广告、推荐模型本身embedding庞大,DNN层逐步向CV,NLP的复杂度看齐。
  • 需要大量的A/B实验,模型需要持续、高频更新上线,导致单个场景需要多份训练和推理任务。
  • 直接面向业务核心场景。每次用户操作(刷视频,新闻,电商,搜索等)都可能触发对应的推荐广告模型。

因为广告、推荐模型在embedding层如此显著的特性,导致这类模型的压缩技术也更有针对性。出了传统模型压缩技术外,这类模型需要特定的技术去解决embedding层在训练和推理阶段的压缩问题。然而广告、推荐模型的压缩技术却不是那么的广为人知。成熟和普及度也相对低一些。本文主要讨论这个领域的技术和应用成果。

特征分析

后面很多优化都需要利用一个关键点:每个特征的重要性是不一样的。
因此,表达每个特征的参数向量的重要性也是不一样的,这里重要性体现在几个方面:

  • 特征的访问频率等统计值
  • 特征对模型效果的影响
  • 特征对应的参数当前的收敛程度

特征效果分析

这是个比较大的话题。推荐广告的特征通常在模型里通过embedding向量学习和表示,所以这里主要讨论如何分析embedding:

  • 简单统计指标。许多简单的统计指标对于近似特征效果有一定作用。比如某个特征值出现的频率,出现在正样本中的频率等。
  • 基于搜索机制。比如shapley等方法,通过排列组合不同的特征集合,并对比不同排列组合下的效果,来得出特征效果。这种方法相对比较暴力。
  • 基于梯度值。这里面方法有不少变种,比如基于梯度和参数的乘积,或者是基于梯度的积分。直观上理解,重要的特征参数对于loss的影响比较大,对应的梯度值也会更突出。

训练时压缩

模型训练主要有两个特殊瓶颈点:

  • 存储,embedding参数GB~TB级,加上优化器参数等,内存有较大的开销。
  • 通信,除了DNN层的通信,embedding层也存在大量通信,按每个sample 1000个embedding向量算,128 batchsize需要10w个向量。而大规模分布式训练导致batch size进一步扩大100倍,达到1000w向量的单步通信量。单向量100~500Byte,总可以超过数GB每步。

存储压缩

动态Swapout

我们发现,特征访问频率存在显著的power-law分布,甚至更极端。超过80%的访问少于20%的特征。因此,基于这个特性,我们可以将低频的特征swapout到磁盘等持久性存储中。实践表明,swapout最低频的50%特征参数不会带来显著的磁盘IO。同时配合无量框架分布式训练是的“prepare机制”,这部分IO可以提前被完成,掩盖掉swapin的延迟。通过这种方式可以在降低显著内存空间的情况下,不影响训练的性能。

其实在小流量推理时,本地的SSD也可能用类似机制。

特征变长和淘汰

基于上面的特征效果分析结论,可以有相应的策略来节省存储:

  • 变长。通常不重要或者说对模型效果影响不大的,我们可以降低对应向量的长度,用更少的参数,或者更低的精度去学习表示,对重要的用更多的参数。
    • 这里有个细节。cardinality比较大的feature field/slot,从信息论角度,似乎也需要更长的向量。
  • 淘汰。当内存预算快用满时,基于特征效果分析,可以优先将最不重要的那部分参数直接扔掉。这种方式对于推荐等具有一定时效性的场景往往简单有效。因为整个推荐系统本身也会清理过时的item,优先推送新热的item。只要阈值控制合理,通常不会对模型效果带来可观测到的效果下降。

其他

其他还要许多在embedding层面思考压缩的方法。比如:

  • 基于特征统计分布的准入,在进入模型前就干掉。这和特征淘汰有一定的功能重叠。
  • 特殊的优化器。这通常是一种组合技术,基于特征效果分析,针对不同特征使用更小开销的优化器(而不是adam),尽量在不影响模型效果的情况下降低优化器参数的开销。

通信压缩

基于参数重要性

前面提到了特征重要性分析。这部分结果也可以简单应用在通信压缩上,比如不重要的特征参数:

  • 可以考虑使用fp16。
  • 可以mask掉大部分接近0的gradient&weight
  • 可以使用更异步的通信机制

特征稀疏化

通信是分布式训练里面一个比较突出的瓶颈,特别是对于基于CPU的分布式训练,2021年通常是基于以太网的25Gbps,而不是nvlink的数百G的带宽。

Facebook提出的稀疏化方案降低通信量达到90%以上,同时甚至能带来微弱的模型效果提升,比如Data Parallel Training时:

  • 每L步根据希望达到的稀疏率,计算需要被通信的w值中最小的值。
  • 小于该阈值的w全部被mask掉,不进行通信。而被mask掉的值会作为error记录下来。
  • error会和下一步的w加和,进行错误补偿。然后w进入新一轮的稀疏化。
  • 这里注意的是“阈值”每L步计算一次,避免所有参数排序开销(压缩、解压开销)

另外还提到Model Parallel Training时的稀疏化技术,甚至还能提升模型的收敛效果。总体来看如下图:

其他

在分布式训练的各个环节都可能存在通信:

  • 跨节点的通信
  • 跨GPU设备的通信
  • CPU和GPU间的通信

以及通信的不同重要性的内容:

  • 参数通信
  • activation通信
  • gradient通信
  • 优化器相关状态通信

对于各个流程数值分布分析,重要性分析和对模型的理解有助于挖掘更多的通信压缩点

---

推理时压缩

推理时压缩的方法非常的多,在推荐广告场景,更多是把传统方法针对领域进行定制化的引入。这里举几个比较有效的例子。

行量化

前面提到推荐广告模型有个巨大的embedding table。每一行的参数向量对应一个特征。传统训练完后是fp32的。我们希望把他量化压缩一下。
量化的方法就非常多了,有uniform, non-uniform。还有就是对outlier的处理策略也非常多样。这里一个关键点是一个向量quantize然后dequantize后和原向量的diff。直观上就是量化的信息损失。

很自然想到的就是为每行的向量分别量化。然后为每一行用一个简单搜索策略来最小化diff。具体的:

  • 计算向量中的min, max。(这里其可以优化一下,减少搜索开销)
  • 不断小步加大min,或者减小max。寻找量化MSE损失最少的min、max。进而得到量化scale。
  • 通过scale值量化向量后保存,同时保存scale。

scale值每一行有一个,本身是个额外开销,所以也可以被压缩:

  • 简单地换成fp16。
  • 用一个长度256的codebook + int8 scale值。可能会带来效果损失。

混合精度量化

在实践中,int8通常能让模型效果不下降。但是对于要求比较高的业务,他们还希望最终的打分分布和原来fp32的结果几乎一样,这就比较麻烦了。另外,有些场景int8也可能带来效果损失。或者当我们进一步量化到int4时可能带来效果损失。

前面提到的特征重要性分析非常重要。可以基于特种重要度分析的结论,做混合精度的量化解决问题:少部分重要embedding进行fp16/int8压缩,其他的用int8/int4压缩。

这里在结合一些机器学习平台提供的自动搜索技术,尽量自动化的完成量化策略选择,保障模型效果损失可接受范围内,最大化压缩模型。

通信压缩

当我们把重要的参数cache在本地时,就可以显著减少访问远程全量的TB级的其他参数。网络的抖动也未必需要重试,规避长尾问题。


协同压缩

如果将视角放大,从整个训练和推理两阶段来看问题,或者从整个召回,粗排,精排的流程来看问题,会提供更多的模型压缩机会,解决局部优化无法解决的问题,比如Post-training量化的效果损失问题。这里提几个点。

粗精排联合训练

在我另一篇关于推荐系统foundation model的文章中提到了这个方法,可以起到模型压缩的效果。具体的:

  • 传统推荐流程,为了防止粗排过滤太多宝贵item,可能会“不必要”的提高粗排复杂度。为什么不必要,因为后面还有个更复杂的精排。
  • 在训练的时候,粗排和精排模型可以放在一个脚本里面同时训练。精排蒸馏粗排可以显著提高两者一致性。
  • 也就是说,我们可以使用更简单的粗排模型达到类似的效果,起到了压缩粗排模型的目的。

训练推理联合压缩

除了传统的training-aware quantization,还有其他手段。其实上面“推理时混合精度量化”稍作改进就可以得出。

  • 在线训练实时的统计特征重要性,并按照特征重要性导出混合精度量化后的模型。
  • 推理阶段可以直接加载量化后的模型。
  • 通过持续在线训练和增量上线,可以不断调整不同embedding的量化值。

进一步探索

前面提到,通常推荐广告模型一个关键特性是99%以上参数是embedding,那么优化embedding是关键。而优化embedding非常依赖特征重要性分析。

其实漏了说,推荐广告模型还有一个关键特性。就是两阶段训练:

  • 离线训练。类似传统CV、NLP模型,从随机初始化到loss收敛阶段。
  • 在线训练。收敛后的模型会持续获得流式数据,不断训练,在分钟级、甚至流式上线新增量。

于是,有几种可能:

  • 在线训练时,高频embedding是否已经收敛,是否可以放松传统Dist-SGD的同步约束,压缩通信。
  • 在线训练时,除了头一两层DNN,后面的DNN是否也可以放松约束,压缩通信。
  • 新特征,老特征,DNN的更新模式是否有差异化更新方式,压缩通信。

作者信息:

Peter潘欣 来源: 推荐广告模型的降本提效:压缩策略


本文地址:https://www.6aiq.com/article/1634257916110
本文版权归作者和AIQ共有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出