训练更灵活、经济、高效，新一代搜索推动稀疏大模型训练范式GBA

2024-11-27 12:28:03 | 来源: 互联网整理

一、搜推广增量迭代挑战

搜索推广模型不仅需要具有高训练加速比的数据恢复能力，还需要能够进行高度并行的定性实验和日常增量训练。稀疏大模型的主流搜索和推广通常采用异步或同步训练模式。在异步训练中，训练通常采用PS（参数服务器）架构，并以高并发的方式执行。每个计算节点所需的资源量并不大，可以利用阿里巴巴内部主机托管策略下的海量资源，适合定性实验。以及每日增量训练。同步训练一般采用AllReduce架构，参数直接放在计算节点上，通过集体通信在节点之间交换。在高性能集群中，同步模式可以充分利用计算和通信硬件资源，从而获得更高的训练性能。理论上可以比同等成本的异步训练任务高很多。

近两年，XDL训练引擎推出的HPC（高性能计算）同步训练在阿里巴巴内部搜索、推荐、广告业务得到广泛推广。相关的开源训练框架HybridBackend也受到了业界的广泛关注。但基于搜索推广模型的迭代特点以及当前的资源状况，既需要HPC模型较高的训练加速比，又需要日常实验迭代的高并发实验能力。如何实现同步与异步训练的自由切换，充分利用HPC高性能资源和海量共置资源，成为提高迭代效率的重要问题。并且为了避免模型迭代过程中在参数调整上浪费时间和资源，保证模型效果稳定，在切换训练模式时需要做好以下几点：

效果：同步模式切换到异步模式后效果不变，获得同步训练的效果；性能：在混合部署的不稳定资源环境下，同步模式切换到异步模式后获得异步训练资源利用效率；鲁棒性：无需切换优化器，无需调整优化器等一系列超参数；可扩展性：可以缓解由于worker数量增加和陈旧性增加而导致的异步训练的影响。但在实践中，我们发现直接切换训练模式并不能达到上述目标，而且实现上述能力也存在很大的挑战：

1）考虑到实际业务对精度损失的容忍度较低，目前业界和学术界对于混合生产环境的训练模式切换还没有成熟的解决方案，因此没有明确的解决路线可循。

2）切换带来的效果问题无法通过简单调整学习率或其他超参数来解决。两种训练模式在批量大小、相同数据下的更新步数以及优化器选择上存在显着差异，这会导致参数更新实际上倾向于不同的收敛点。切换后的短期精度损失（甚至直接训练崩溃）以及每个训练阶段结束时潜在的精度不一致是日常迭代开发中无法容忍的。

从高效迭代的角度来看，需要一种可以在高性能同步训练和高并发异步训练之间自由切换的训练模式，而不需要用户中途调整超参数。桥接两种训练模式的核心问题是如何在训练时使两种模式的梯度分布大致相同，从而使模型向相同或相似的收敛方向更新。我们观察到梯度分布和梯度延迟是影响效果的重要突破口。基于此，我们提出了一种基于全局梯度聚合（GBA）的训练模式，首先累积几个梯度，然后进行梯度更新，直到总批量大小相同。实验证明GBA和同步训练可以自由切换，无需超参调整。

1.1 业界相关工作

近年来，业界一直在探索介于同步和异步之间的“半同步”训练模式，以避免慢节点的影响，同时避免效果损失。例如，Tensorflow中集成的Backup Worker方法每次同步更新时都会自动忽略指定数量的慢速节点的梯度；另一个例子是XDL中限制节点之间梯度版本的原始方法，它允许有限的梯度过期。同时，还有一些与“半异步”训练模式相关的工作，比如每次更新几个batch的梯度，或者完全去相关稀疏参数的获取和更新，这些都是有效缓解的方法。慢节点的影响。然而，虽然这些训练模型在开源数据集上都报出了理想的结果，但在实际的工业生产场景中，数据丢失和严重的梯度过期对业务精度的影响是不可接受的，而快慢节点速度的巨大差异就是这样的。这是普遍的情况，所以要想真正落地，现有的这些学术设计还远远不够。

二、GBA的诞生

2.1 源于生产实验的观测与启示

2.1.1 直接切换训练模式后，效果出现严重损失

图1 切换训练模式后AUC变化趋势

我们将DeepFM 和Twin Towers 模型的训练任务分布在Criteo 竞赛数据集和生产数据集上。最初，我们选择同步或异步训练，在消耗一半数据集后，我们切换到另一种训练模式继续训练。图1绘制了不同训练阶段的测试AUC，每训练进度的5%进行一次测试。可以看到，刚发生训练模式切换后，无论切换到哪个方向，都会有非常明显的效果损失，然后效果会慢慢恢复，或者AUC下降到0.5并且无法恢复。这表明，在没有特殊设计或精细参数调整的情况下，切换训练模式将对模型效果产生灾难性影响。

图2 双塔模型未切换和切换后一层参数数值分布变化

进一步观察发现，异步和同步训练本质上会导致模型在不同的收敛方向上更新。如图2所示，异步和同步训练后某一层的参数分布完全不同。当异步训练切换到同步训练一段时间后，参数分布被强制训练并更新为同步分布。这意味着，为了获得平滑切换，确保参数分布保持不变是一个值得考虑的出发点。

2.1.2 梯度的分布与更新梯度时实际的batch大小有关

图3 异步梯度聚合方法和同步训练方法的梯度数值分布

参数分布相似的前提是梯度分布也相似。我们使用异步梯度聚合进行了大量实验，发现当batch size与同步训练的全局batch size相同时，两种训练模式的梯度分布会比较接近（如图3所示）。考虑到集群整体硬件资源紧张程度的变化，一天中不同时间的执行结果略有不同。资源紧张会带来不同计算节点执行速度的较大差异，从而导致平均梯度过期较高。我们发现在两个实验中梯度过期的程度与梯度分布的偏差呈正相关。这一发现启发了使用相同大小的批次进行梯度更新并确保梯度过期较小，这是实现平滑切换的潜在解决方案。

2.2 全局梯度聚合机制

基于上述实验分析，我们引入全局批量大小的概念，其定义为实际梯度更新聚合的样本大小。我们将Local batch size（即每个worker处理的batch size的大小）和同步训练模式下的worker数量分别定义为和。同步训练中的全局批量大小定义为和。根据分析和观察的结论，当我们切换训练模式时，GBA将保持Global batch size不变。 XDL基于PS架构，设计了全局梯度聚合机制。在这种机制下，每个计算节点消耗数据并计算梯度，然后将梯度发送到PS节点。梯度不会立即应用，而是在累积到与同步的全局批量大小一致时进行聚合以执行参数更新。因此，计算节点不会被梯度更新阻塞，快的节点可以消耗更多的数据，慢的节点只需要消耗更少的数据，计算能力得到充分调动。

对于每个更新步骤，所有密集层参数都将被更新，但只有一小部分稀疏层参数（即嵌入参数）被更新。两个参数的梯度过期程度不一致，因此很难用参数版本来控制。我们这里统一将梯度过期程度定义为训练过程中数据的过期程度。数据的陈旧程度表示工作人员开始训练其本地批次时的全局步骤与执行梯度更新时的全局步骤之间的差异。在同步训练过程中，数据的陈旧度始终为0。基于数据陈旧度，我们在PS架构上实现了令牌控制的GBA机制，以实现梯度聚合，并处理异步训练时稀疏模型的梯度过期问题。

图4 GBA全局梯度聚合机制

图4展示了GBA中核心的Token控制机制。在传统的PS 角色上，我们设置了一个名为Token List 的队列，其中按顺序包含Token。队列中每个Token的初始编号为Worker。 Worker每次请求参数时，都会获得一个Token。计算梯度后，所有Token和反向相关数据将被发送回PS。 PS上还设置了一个Gradient Buffer，用户缓存梯度而不立即更新。每当累积的batch size与同步的全局batch size一致时，PS就会聚合Gradient Buffer中的梯度并更新参数。 Gradient Buffer不会关注梯度来自哪个Worker，而是会关注Token版本。如果Buffer 中的某个Token 落后于当前Token List 已发行的最新Token 超过给定的阈值，那么我们认为梯度已经严重过期，需要被丢弃，而不是参与梯度更新。这样GBA既保证了实际更新梯度时的batch size，又避免了节点速度慢和梯度严重过期带来的不利影响。

论文中，我们结合稀疏参数和稠密参数不同梯度延迟敏感度的特点，对GBA进行了偏差分析。结合多次训练步骤后的误差期望，我们证明GBA训练模式不会影响模型的收敛。

三、GBA收敛性分析

业界关于同步异步训练收敛性分析的工作有很多。我们参考了一些工作，将同步的收敛公式形式化为训练步骤后的误差期望：

其中，表示步长的参数，表示学习率，表示Lipschitz常数，表示梯度的标准差。是一个经验风险函数，并假设具有强凸系数。是训练结果与最优解之间的差距。公式中的第一项代表误差的上限，代表误差的衰减率。

我们提出的GBA是一个基于异步梯度聚合的方案。我们假设当存在并且

这里，是梯度过期对参数更新影响的衡量因子。该因子越小，梯度过期对模型效果的影响越小。另外，它表示参数在token索引和token索引处同时更新的概率。对于稀疏模型，这将是一个远小于1 的数字，而对于密集参数，该值为1。GBA 的误差公式可以表示为：

其中， 和是代币索引等于全局步长的概率下界，即非梯度过期的概率下界。考虑到同步和异步收敛公式，如果设置为相似的值，GBA可以获得与同步相似的收敛特性。这也与我们GBA中设定的全球批量大小方案一致。特别是对于比CV或NLP模型更小的稀疏模型，在稀疏模型场景下使用GBA也会获得并同步更小的收敛差异。

四、GBA工程实现

图5 GBA项目实施逻辑

GBA整体工程实现架构如图5所示。 PS端主要包括三部分逻辑：生成token并根据token控制过时（黄色部分逻辑）、累加梯度并异步聚合更新参数（绿色部分逻辑）、无需等待即可响应worker请求（蓝色部分逻辑）逻辑）。 Worker端的主要流程（红色部分逻辑）是每次请求参数和token（根据TF设计，训练的第0步不需要请求token）。节点中计算出梯度后，上传到PS，无需等待即可开始下载。一个批处理过程。

1、如何保证GBA达到同步效果和异步速度？

首先，GBA的梯度聚合不需要梯度来自不同的worker。 GBA将同一步的梯度定义为任何worker在两次参数更新之间发送的梯度。系统将在每一步中发行worker_num代币。此步骤中的所有代币都可以由任何工人获得。即使存在超慢节点，这一步的所有代币仍然可以被拿走，触发下一步的代币分配。因此，只要有worker运行任务，代币的生产和消费就可以正常进行。 GBA利用ps的梯度聚合操作，保证梯度根据全局batch更新，获得同步效果；它通过执行任务来获得异步速度，而无需工作人员等待。

2. GBA模式会丢失很多梯度吗？速度和数据完整性之间如何权衡？

因为令牌队列是优先级队列，任何worker获得的令牌都是当时队列中最小的令牌。假设慢速worker在获取sync_step后比其他worker慢10步。当慢速worker再次获得令牌时，它会获得sync_step + 10，因为[sync_step + 1,sync + 9]已被其他worker运行。因此，只有当该慢速工作线程比其他快速工作线程慢许多倍时，才会发生梯度丢弃。但既然这个慢worker的梯度已经严重过期，丢弃梯度也不是什么坏事。而且对于每天使用大量样本的搜索和推广模型，丢弃少量梯度不会影响效果。

五、GBA实验效果

图6 GBA与典型半同步/半异步训练模式和同步训练切换后AUC变化趋势

我们在Criteo-1TB数据集上实现了DeepFM，在阿里妈妈开源CTR数据集上实现了DIEN，并在生产数据集上实现了简化的生产双塔模型，采用GBA等半同步/半异步方法。训练，并进行同步训练的双向切换，观察GBA的实际表现。图6显示了切换发生后AUC的变化趋势。与保持同步训练相比，GBA的总体趋势是相同的，而其他方法在至少一项任务上表现出明显的下降，特别是当切换到自同步时，其他训练模式都有明显的“重新收敛”过程。从数据来看，GBA在继承后第一天、数据集最后一天、多日平均值三个指标上都取得了最好的表现，实现了同步训练之间的平滑切换。

阿里巴巴集团内部，多个广告业务场景已经使用GBA/同步切换的方式进行实验和迭代。在这些商业模式上，同步训练和GBA模式训练的线下效果是一致的。将训练范式从异步改为GBA/同步切换后，我们观察到部分模型上千分位数的AUC有所提升，这对于提升搜索推广模型的线下和线上效果非常显着；同时，使用继承参数继续训练已经不是唯一的选择，用过去一年的数据从头开始训练只需要两周的时间，这使得GBA/同步切换成为非常流行的解决方案。

六、未来展望

GBA/同步切换的训练范式下，当模型版本发生变化，需要大量数据进行训练时，利用高性能HPC资源进行训练，尽快获得实验结果和模型参数；而在日常日级迭代或者ODL训练过程中，可以切换到GBA模式，使用普通计算节点执行训练和交付任务。尤其是对于创新探索场景，HPC同步+GBA的组合模式可以帮助算法工程师摆脱对“祖传参数”的强烈依赖，最大化探索空间。目前，XDL也在设计自动迁移训练模型的解决方案，希望成为标准化、通用的解决方案，服务于更大规模的搜索和推广稀疏模型训练。欢迎对这项工作、搜索推广算法迭代、训练框架等感兴趣的研究人员和工程师与我们交流。我们非常期待与您共同探讨相关领域未来的发展方向，并产生更多有价值的实施成果。