BIGO技术 | 万亿模型参数的训练

*本文转载自 BIGO技术 公众号，对BIGO技术感兴趣的同学可以关注了解更多内容~

 

​一、推荐场景特征稀疏性

推荐场景通常由于引入了大量的ID类特征从而导致存在海量稀疏参数，例如下图经典YouTube DNN模型中，使用用户观看过的视频以及用户历史search tokens作为主要Embedded特征。根据论文中论述，YouTube DNN中candidate video以及search tokens均有百万之巨。在此基础上如果再使用交叉特征，就会使参数爆炸问题进一步加剧。

图1  YouTube DNN

 

在BIGO推荐场景主模型的稀疏参数也已经近万亿，即使按照FP16存储也要消耗近1T的内存空间。BIGO机器学习平台团队针对推荐场景建设了大规模稀疏参数的训练系统，通过参数服务器的深度优化，双图训练，精细化的特征准入和退出策略，自动统计特征，异构硬件平台下的在线预估性能调优，在模型训练，模型交付，在线预估缓解下完整的解决了参数维度爆炸问题，本文主要介绍我们的一些实践和技术思考。

 

二、参数服务器

万亿级参数的模型训练首先要解决参数通信问题，Ring-AllReduce和独立参数服务器是业界两种主流的参数通信模式。

Ring-AllReduce方式将Parameter Server和Worker合为一个实例。训练过程中，每个实例直接利用本地完整参数副本进行Forward和Backprop计算，有以下特点：

1、单个实例可存储完整参数；

2、要求采用参数同步更新机制，适合对训练速度要求不高但每次Mini-Batch梯度需高精度更新的场景；

3、在N个设备，数据总量K情况下。每个设备在Scatter Reduce阶段，接收N-1次数据，在All Gather 阶段，接收N-1次数据。每个设备每次发送 K/N 大小数据块。故总数据传输量=2(N−1)*K/N，随着设备数量 N 增加，总传输量为恒定。

图2  Ring-AllReduce过程示意图

 

参数服务器（Parameter Server）将参数的存储和更新上升为一个独立的角色，即在分布式调度时Parameter Server与Worker角色分别对应为不同的实例，与Ring-AllReduce相对应有如下特点：

1、支持海量稀疏参数。整体参数经过Sharding后分别存储在不同Parameter Server实例上；

2、支持同步和异步两种参数更新方式；

3、每次Mini-Batch开始时，各个Worker从Parameter Server拉取本次迭