苗旭鹏¹, 张敏旭¹, 邵蓥侠², 崔斌¹
1. 北京大学 信息科学技术学院, 北京 100871;
2. 北京邮电大学 计算机学院, 北京 100876
收稿日期：2021-07-22
基金项目：国家重点研发计划项目(2018YFB1004403); 国家自然科学基金项目(61832001); 北京大学-腾讯协同创新实验室项目
作者简介：苗旭鹏(1995—), 男, 博士研究生
通讯作者：崔斌, 教授, E-mail：bin.cui@pku.edu.cn

摘要：传统的分布式深度学习训练系统大多基于参数服务器和全局规约通信框架, 缺陷日益显著：参数量大, 基于全局规约的去中心化通信架构由于无法存储全量模型而无法使用; 通信量大, 基于参数服务器的中心化通信架构面临着严重的通信瓶颈。为了解决以上问题, 该文提出了面向大规模深度学习推荐模型的混合通信训练框架PS-Hybrid, 分离了嵌入层参数和其他参数的通信逻辑, 实现了PS-Hybrid原型系统。实验结果证明了所提出的混合通信方案能够比纯参数服务器方案取得更好的性能, 在16个计算节点下比TensorFlow-PS加速48%。
关键词：推荐模型分布式深度学习参数服务器全局规约
PS-Hybrid: Hybrid communication framework for large recommendation model training
MIAO Xupeng¹, ZHANG Minxu¹, SHAO Yingxia², CUI Bin¹
1. School of Electronics Engineering and Computer Science, Peking University, Beijing 100871, China;
2. School of Computer Science, Beijing University of Posts and Telecommunications, Beijing 100876, China

Abstract: Most traditional distributed deep learning training systems have been based on parameter servers which have centralized communication architectures that face serious communication bottlenecks due to the large amounts of communications and AllReduce communication frameworks which have decentralized communication architectures that cannot store the entire model due to the large number of parameters. This paper presents PS-Hybrid, a hybrid communication framework, for large deep learning recommendation model training which decouples the communication logic from the embedded parameters and other parameters. Tests show that this prototype system achieves better performance than previous parameter servers for recommendation model training. The system is 48% faster than TensorFlow-PS with 16 computing nodes.
Key words: recommendation modeldistributed deep learningparameter serverAllReduce
近年来，深度学习技术在多个领域得到了广泛应用，取得了显著效果。深度学习模型的参数数量通常决定了模型的表达能力上限。为了取得更好的模型性能，模型的参数规模变得越来越大，例如谷歌最新发表的Switch Transforme^[1]的参数规模达到了1.6万亿个。在实际工作中，以最常见的推荐问题为例，由于具有海量高维的输入特征，其模型参数量通常可以达到千亿甚至万亿量级。与传统的深度学习模型不同的是，这种大模型的参数主要以嵌入表的形式组织。嵌入表通常用来建模类别特征，表的每一行都是一条固定维度的嵌入向量，用来表示一个类别特征。工业领域常见的大规模深度学习推荐模型如深度线性网络(wide & deep network)^[2]、深度交叉网络(deep & cross network，DCN)^[3]等都采用了这一结构。
由于工业领域模型特征维度很高，导致整体模型的参数量巨大，这种大规模深度学习模型往往需要分布式训练。目前，主流的分布式深度学习系统主要采取2种通信架构：1) 参数服务器(parameter server，PS)架构^[4-6]。典型的系统包括TensorFlow^[7]等，模型参数被集中存储到参数服务器集群，计算节点通过与服务器进行交互来完成模型参数的拉取和更新，通信方式灵活，可以异步化。2) 全局规约(AllReduce)架构。典型的系统包括PyTorch^[8-9]等，每个计算节点存有一份完整的模型参数副本，通过节点间的高带宽网络以去中心化的方式进行模型参数的同步，由于通信模式固定，通常只能采用同步方式进行通信。然而，对于大规模的深度学习推荐模型，现有方案仍然面临以下挑战：
1) 模型规模大，全局规约等去中心化方案无法存储全量的模型参数；
2) 通信量大，中心化的参数服务器通信架构无法充分利用显卡(GPU)间的网络带宽。
针对以上两个问题，本文提出了面向大规模深度学习推荐模型的混合通信架构。具体来说，本文首先分析了此类模型中嵌入层的稀疏访问特性，基于该特性，提出了混合通信方案。具体地，嵌入层的稀疏模型参数部分使用参数服务器来进行同步，而稠密参数部分使用全局规约来同步，通过合理利用GPU间的带宽来减少通信开销，从而在不影响模型收敛速度的前提下，大幅缩短推荐模型的训练时间。另外，针对传统参数服务器通信机制中迭代锁导致的空闲等待，提出了基于优先级传输的优化方案，从而加快下一轮迭代的开始时间，使得计算和通信充分重叠。为了验证该混合通信方案，本文基于经典的参数服务器PS-Lite^[10]和集合通信库(NVIDIA collective communication library, NCCL)^[11]实现了混合通信系统PS-Hybrid，并且在开放数据集Criteo和Avazu上进行了充分的实验。实验结果表明：基于混合通信的推荐模型训练，在16个计算节点数据并行训练的设置条件下，PS-Hybrid比TensorFlow-PS加速48%，且具备良好的扩展性。
本文的贡献可以概括如下：
1) 对大规模深度学习模型训练的不同网络通信架构进行了定量分析，基于参数服务器和全局规约通信架构，提出了混合通信框架用于提升此类模型的训练性能；
2) 对传统的参数服务器方案中存在的迭代锁进行了优化，通过梯度的优先级传输提升计算和通信的重叠度；
3) 实现了基于PS-Lite和NCCL的混合通信方案，并且用充分的实验证明了在推荐模型训练场景中，混合通信方案比纯参数服务器方案更高效。
1 相关工作1.1 深度学习推荐模型深度学习推荐模型通常采用先低维嵌入再接多层感知机的结构，其模型训练流程如图 1所示。训练分为2个主要模块：嵌入/叠加层和特征交互层。嵌入层参数就是嵌入矩阵，特征交互层主要是多层感知机。首先处理样本，对其中的类别特征进行编码，转换为嵌入表的索引；然后根据索引进行嵌入表的查找操作，样本的每个类别特征被转换为一个低维稠密向量。叠加层将样本的类别特征和数值特征进行拼接，得到包含样本所有特征的特征向量并将它输入到特征交互层，特征交互层对特征向量进行复杂的交互操作。最后，使用特征交互层的输出进行分类预估。

图 1 推荐模型训练流程

图选项

1.2 参数服务器和全局规约训练数据量的快速增长使得单个GPU很难在有限时间内快速训练好模型。常见的分布式训练框架通常采用数据并行的分布式训练策略，具体包括2种通信架构：参数服务器和全局规约。
参数服务器^[12]架构由计算节点和服务器节点组成，其中服务器节点以键值对形式集中存储全局共享的模型参数，模型参数通常进行分片并均匀地存储在服务器集群中。在每轮迭代中，计算节点首先基于本地数据计算得到模型参数的梯度并发送给服务器节点，服务器节点接收来自计算节点的梯度并对模型参数进行更新，计算节点拉取最新的模型参数用于下一轮的迭代。计算节点间的参数同步通过服务器节点间接完成，能够方便地支持容灾且提供完全同步、完全异步、半同步3种模型同步策略，较为灵活。同步训练方式严格等价于单机的模型训练，模型收敛得到保证，但是因为每轮迭代都要求计算节点互相等待，整个训练任务会被最慢的计算节点拖累；异步训练中，每个计算节点不会等待其他节点，但会遇到梯度陈旧性问题，导致模型收敛质量变差。由于计算节点不需要互相等待，因此迭代速度较快。
另一种数据并行的训练框架是全局规约框架。它是一种集合通信操作，将所有计算节点中的同一个参数向量进行规约操作，得到一个规约后的参数向量，其中的每个元素是所有参数向量对应位置上元素的和或者平均值，随后将该参数向量再广播回所有计算节点。以全局规约架构进行数据并行训练时，在每轮迭代中，每个计算节点基于本地数据计算出梯度后，调用全局规约操作对梯度进行聚合后更新本地模型参数。全局规约有多种实现方式，常见的有树规约^[13]、环形规约^[14]和层次规约^[15]。在环形规约架构中，每个计算节点只与邻居节点通信，并且保证每个计算节点用于同步梯度的数据通信量是常量，与计算节点个数无关。由于具有相对固定的通信架构，全局规约架构只能以同步的方式进行通信，不如参数服务器架构灵活，容错也较为困难。
随着计算节点数量的增加，参数服务器节点会成为整个系统的通信瓶颈，需要通过参数的合理分配来避免节点的热点问题，而全局规约架构中每个计算节点的通信量与计算节点的数量无关，因此不会遇到瓶颈问题。
2 基于参数服务器的推荐模型训练2.1 推荐模型参数组成分析从图 1可知，推荐模型中的参数分为两类：嵌入矩阵和稠密参数。这两类参数在存储和更新上有鲜明的区别：从存储的角度考虑，嵌入矩阵的大小由所有类别特征的类别数总和以及嵌入向量的维度决定，参数数量能够达到千亿级别。在真实场景中，单机不能存储下如此庞大的参数，需要多机分片存储。稠密参数由于规模较小，可以完整存放在单个GPU中。从参数更新角度考虑，嵌入矩阵每一轮迭代只更新部分嵌入向量，取决于当前批的训练数据，稠密参数在每一轮迭代中，所有的元素都会被更新。参数的存储特性和更新模式决定了模型训练框架的选择。考虑到存储特性和模型按需获取嵌入向量，在实际场景中，训练推荐模型都采用参数服务器的训练框架，而全局规约框架要求单个计算节点保留有完整模型参数，因此无法使用。
2.2 基于参数服务器的推荐模型训练方案基于参数服务器的推荐模型训练流程如图 2所示。

图 2 基于参数服务器的模型训练

图选项

训练流程如下：
步骤1??所有模型参数在参数服务器端初始化。嵌入矩阵按照索引维度均匀切分并分片存储，稠密参数按照向量维度均匀切分并放置到参数服务器上。
步骤2??每一轮迭代中，每个计算节点遍历本地输入的批数据，对于类别特征，获取去重后的嵌入矩阵索引，从参数服务器拉取索引对应的嵌入向量以及完整的稠密参数。
步骤3??拉取最新参数后，模型进行前向和反向计算。在计算出梯度后，对于稠密参数的梯度，调用参数服务器的稠密通信接口，对于嵌入矩阵的梯度，调用参数服务器的稀疏通信接口，将梯度发送回参数服务器并拉取更新后的参数。
3 基于混合通信的推荐模型训练下面首先从网络通信的角度分析参数服务器方案的不足以及全局规约方案的优势，进而提出了基于混合通信框架的训练方案。
3.1 通信量分析推荐模型的参数主要分为嵌入矩阵和稠密参数。本文对两类参数分别使用参数服务器架构和全局规约架构进行一轮训练迭代的通信量进行分析。在参数服务器架构中，本文假设分别有N个参数服务器和N个计算节点。在全局规约架构中，本文假设有N个计算节点。
3.1.1 稀疏嵌入参数通信量分析首先分析稀疏嵌入参数在不同通信架构中的通信量。假设有M个稀疏参数，参数的平均大小为P字节；且每一轮迭代，每个计算节点只访问每个稀疏参数的一部分元素，本文假设这个比例为α。
参数服务器?? 本文假设稀疏嵌入参数被均匀切分到各个参数服务器。在每一轮迭代中，对于每个计算节点，它拉取αMP字节，发送αMP字节，N个计算节点的通信量就是2αMNP字节；对于每个参数服务器，它接收αMP字节，发送αMP字节，通信量为2αMP字节，N个参数服务器的通信量就是2αMNP字节。综上，对于M个稀疏参数，参数服务器方案总通信量为4αMNP字节。
全局规约?? 本文采用环形规约。对于M个稀疏参数，环形规约退化成了N-1轮的All-Gather迭代，每个计算节点的通信量为2αM(N-1)P字节，N个计算节点的通信量为2αMN(N-1)P字节。
根据以上分析，M个稀疏参数在参数服务器和全局约架构下的通信量分别为4αMNP和2αMN·(N-1)P字节。当满足N-1≥2的条件时，使用全局规约架构来训练稀疏参数比参数服务器架构的通信量要大，且该条件在大规模的训练集群中是肯定满足的。此外，全局规约在每个计算节点上都需要保存完整的稀疏参数，需要更大的内存资源，因此稀疏参数更适合使用参数服务器架构来训练。
3.1.2 稠密参数通信量分析假设有M个稠密参数，参数平均大小为P字节。
参数服务器?? 稠密参数按照向量维度切分，M个参数均匀分布在N个参数服务器上，每个参数服务器维护M/N个稠密参数。在每一轮迭代中，对于每个计算节点，它拉取M个参数、发送M个参数，通信量为2MP字节，N个计算节点就是2MNP字节；对于每个参数服务器，因为维护M/N个参数，通信量为2MP字节，那么N个参数服务器通信量就是2MNP字节。综上，参数服务器总通信量为4MNP字节。
全局规约 ?? 以环形规约为例，一次同步包括Reduce-Scatter和All-Gather两个N-1轮的迭代。每个计算节点的通信量为4M(N-1)P/N字节，N个计算节点的通信量为4M(N-1)P字节。
根据以上分析(表 1)，M个稀疏参数在参数服务器和全局规约架构下的通信量分别为4αMNP和2αMN(N-1)P字节。当满足N-1>2的条件时，使用全局规约架构来训练稀疏参数比参数服务器架构的通信量要大，且该条件在大规模的训练集群中是肯定满足的。训练集群规模越大，通信量差距越大，例如当N=33时，全局规约架构的通信量是参数服务器的8倍。此外，全局规约在每个计算节点上都需要保存完整的稀疏参数，需要更大的内存资源，在稀疏参数规模较大的情况下是无法满足的，因此稀疏参数更适合使用参数服务器架构来训练。M个稠密参数在参数服务器和全局规约架构下的通信量分别为4MNP和4M(N-1)P字节，因此如果将稠密参数通过环形规约架构来同步，比参数服务器可以节省4MP字节的通信量，并且稠密参数量越大，使用环形规约可以节省的通信量越大。
表 1 通信量分析

参数类型	通信架构	数据通信量/字节
稠密	参数服务器 全局规约	4MNP 4M(N-1)P
稀疏	参数服务器 全局规约	4αMNP 2αN(N-1)MP

表选项






从网络传输角度考虑，如果稠密参数通过全局规约来同步，还可以充分利用GPU间的高带宽，尤其在单机多卡的设置下，可以消除低带宽的跨机器网络通信。主流的集合通信库NCCL提供了高度优化的集合通信实现，支持GPU Direct P2P和共享内存通信。稠密参数量的规模越大，NCCL的性能优势越能体现出来。
综合以上分析可以发现，从通信角度，稀疏参数和稠密参数分别适合使用参数服务器和全局规约来同步。由2.1节推荐模型的参数组成分析可知，推荐模型同时具备这两类参数。为了进一步节省通信量，同时能够充分地利用GPU间的带宽资源来提升系统整体的资源利用率，本文引入基于混合通信的推荐模型训练方案。
3.2 基于混合通信的推荐模型训练方案基于混合通信的推荐模型训练流程如图 3所示。

图 3 基于混合通信的模型训练

图选项

3.2.1 训练流程步骤1??嵌入矩阵在参数服务器端初始化，稠密参数在每个计算节点本地初始化。嵌入矩阵按照索引维度均匀切分并分片存储。稠密参数在每个计算节点都存有完整的一份。
步骤2??迭代开始，每个计算节点遍历输入的本地小批量数据，获取当前批数据的所有类别特征中出现的嵌入索引，然后从参数服务器拉取索引对应的嵌入向量。稠密参数在本地，不用拉取。
步骤3??模型进行前向和反向计算。为了减少因通信带来的延迟，每个参数在计算出梯度后立即进行异步通信，使得模型计算和通信充分重叠。对于嵌入矩阵参数的梯度使用参数服务器的稀疏通信，对于稠密参数的梯度使用NCCL通信。
3.2.2 系统设计与实现在现有工作的基础上，基于深度学习训练系统河图(Hetu)^[16]实现了混合通信方案。该系统的分布式训练模块改造自经典的参数服务器开源系统PS-Lite，提供了基于内存的Hash映射存储、节点同步、参数分区、参数推送和拉取等基础功能。虽然该开源系统被广泛应用，但是只能提供最基础的分布式训练功能，需要进行定制化的开发和系统优化才能应用于实际场景的模型训练。为了支持大规模推荐模型的训练，本文在PS-Lite的基础上主要添加了如下功能：1) 针对稀疏参数的更新和拉取接口；2) 高效的线程池实现。经过以上优化后，该深度学习系统已经能够对常见的神经网络模型进行高效训练。
河图深度学习系统基于静态计算图的理念进行构建。本文封装了PS-Lite的通信接口供该系统调用。河图系统通过在梯度算子后接入通信算子，并在通信算子中调用封装的通信接口来实现分布式训练功能。该设计的优点在于与系统本身耦合程度低，系统可以接入多种通信算子，而不用考虑算子本身的实现方式。基于该设计，本文系统接入了参数服务器、全局规约、混合通信3种分布式通信框架，具备完善的分布式训练能力且具有较高的训练性能。
为了实现混合通信方案，本文作了如下改进：
1) 在PS-Lite系统里集成了NCCL通信算子。如果使用混合通信训练模式，则在启动参数服务器架构中的计算节点进程时，额外调用NCCL初始化接口，使所有计算节点构成一个逻辑环状的通信集合，用来同步模型中的稠密参数。
2) 增加IndexedSlices类用来专门表示稀疏嵌入表的梯度类型。在混合通信模式中，在梯度通信前对梯度类型进行判断，如果是IndexedSlices类就表示该梯度是稀疏嵌入表的梯度，系统调用参数服务器的通信接口，否则就是稠密参数，系统调用全局规约接口。
参数服务器和计算节点都是运行在服务器上的进程。参数服务器的通信使用PS-Lite的ZeroMQ库，该通信库是对套接字(socket)通信的封装。本文的系统实现没有针对NCCL通信本身作任何修改，NCCL会根据计算节点的分布自动构造最优的通信拓扑。每个计算节点是运行于一个GPU上的进程。当计算节点处于同一个机器上时，会尝试使用GPU Direct P2P和共享内存进行通信。当计算节点处于不同的机器上时，根据硬件是否支持来决定使用远程直接数据读取(remote direct memory access, RDMA)通信或者socket通信。本文实验中的跨机通信全部为socket通信，因此跨机的GPU通信是全局规约通信的瓶颈。
3.3 网络通信优化参数服务器在网络通信上仍然有可优化的空间。一般稀疏嵌入矩阵的查表操作先于稠密参数执行，在反向计算梯度时，稠密参数的梯度会先进入发送队列，然后是嵌入矩阵的梯度。如果用先进先出(first input first output, FIFO)的通信策略^[17]，那么稠密参数的梯度会阻塞嵌入矩阵梯度的发送，推迟下一轮迭代的开始时间。
对于参数在当前迭代反向计算出梯度并进行通信同步到下一轮迭代被用于前向计算的间隔，越靠近输出层的模型参数间隔越大而越靠近输入层的模型参数间隔越小。本文根据这种间隔差异赋予参数的梯度不同的优先级。嵌入矩阵的梯度应该具有最高的优先级。当优先级更高的梯度进入发送队列时，它应该抢占优先级较低的梯度的位置。由于一份梯度无法在发送中途停止，本文提出对梯度分片的方法来实现细粒度抢占优先通信，分片粒度是一个可调整的超参数。
如图 4所示，如果采用梯度优先级传输，梯度通信不仅能和反向传播计算重叠，甚至能和下一轮的前向传播计算重叠，相比于没有应用优先级传输的策略，下一轮训练迭代提前了t′的时间。这种优化方法进一步提升了模型训练中计算和网络通信的重叠。

图 4 梯度切分和优先级传输

图选项

4 实验与分析4.1 实验设置4.1.1 模型与数据集笔者采用3个被广泛应用的推荐模型Wide & Deep^[2]、DCN^[3]和Deep Crossing^[18]，以及2个公开的点击率预估数据集，即Criteo和Avazu数据集，来评估本系统。模型的基本信息概括在表 2中。
表 2 模型参数规模

模型	稀疏参数量/个	稠密参数量/个
Wide & Deep	4 321 609 856	137 984
DCN	4 321 609 856	1 009 755
Deep Crossing	4 321 609 856	111 659 561

表选项






模型的稀疏参数量取决于数据集的所有类别特征的类别数总和，表 2仅列出Criteo数据集上模型的稀疏参数量。经统计，Criteo数据集中所有类别特征的类别总数为33 762 577，Avazu数据集中所有类别特征的类别总数为1 544 258。对于Criteo数据集，本文选取前500万个样本作为训练集。
4.1.2 实验环境实验环境包含4台相同的物理机，每台的配置见表 3。
表 3 实验环境配置

操作系统	Ubuntu 16.04.5 LTS
处理器型号	Intel(R) Xeon(R) Gold 5120 CPU @ 2.20 GHz
处理器个数	56
GPU型号	NVIDIA GeForce Titan XP 24 GB
GPU个数	8
内存	376 GB
网络带宽	1 Gb/s

表选项






4.1.3 基线系统与衡量指标笔者基于TensorFlow1.14、NCCL和PS-Lite分别实现了3种模型，并对比了TensorFlow-PS (TF-PS)、PS-Lite、PS-Hybrid 3种训练框架在相同环境和负载下的各项指标，包括模型收敛性、训练时间、框架扩展性等。
4.1.4 实验设置笔者从以下3个角度设计实验：1) 选用3种常见的推荐模型以及2个公开的点击率预估数据集，以证明PS-Lite和PS-Hybrid针对推荐模型训练场景的通用性和高效性；2) 对推荐模型在参数服务器和混合通信的2种分布式训练框架下的模型收敛情况、训练时间和扩展性进行分析，以证明混合通信方案不影响算法的正常收敛并能大幅缩减训练时间，具有良好的可扩展性；3) 通过调整模型的稠密参数规模来验证混合通信方案在稠密参数规模较大时的优势。在所有实验中，嵌入矩阵向量维度均为128维，均使用随机梯度下降算法进行训练。
本文采用分离式的参数服务器架构：共4台物理机，2个服务器节点分布在其中2台不同的物理机上，若干计算节点均匀分布在另外2台物理机上。
4.2 PS-Hybrid实验结果与分析在PS-Hybrid实验中，本文主要对比混合通信框架与纯参数服务器框架在模型收敛性、训练时间以及框架扩展性等方面的差别。
4.2.1 模型收敛性本文在Criteo数据集和Avazu数据集上分别选取Wide & Deep和DCN模型来展示不同训练框架的模型收敛性，结果如图 5和6所示。本文主要关注模型训练过程中的测试集损失(Loss)和准确度(area under curve，AUC) 2个指标的变化，数据在8个计算节点的设置下获得。

图 5 Criteo数据集上不同框架Wide & Deep模型收敛性

图选项

图 6 Avazu数据集上不同框架DCN模型收敛性

图选项

随着训练的进行，基于不同训练框架的Loss都在稳定下降，AUC在稳定提升。其中：在Loss曲线上，PS-Hybrid比TF-PS和PS-Lite都下降得更快；在AUC曲线上，PS-Hybrid比TF-PS和PS-Lite上升得更快。这是因为引入全局规约的PS-Hybrid在稠密参数更新上使用的是同步方式，而PS-Lite和TF-PS都是异步方式，所以PS-Hybrid模型收敛更快，体现出混合通信方案在训练模型收敛性上的优势。
4.2.2 训练时间比较图 7展示了在相同实验环境下3种训练框架在不同模型上每个轮次(epoch)训练时间对比。计算节点数设置为8，实验数据集为Criteo。DCN模型的交叉层数设为5；Wide & Deep模型的Deep网络隐藏层维度设为1 024，隐藏层数为3；Deep Crossing模型的隐藏层维度设为128，隐藏层数为3。

图 7 不同训练框架的训练时间比较

图选项

图 7表明，在DCN模型上，PS-Hybrid比PS-Lite和TF-PS大幅缩短了训练时间。具体来看，PS-Hybrid比PS-Lite节省了9%的训练时间，比TF-PS则节省了32.5%的时间；在Wide & Deep模型上，PS-Hybrid比PS-Lite节省了37%的训练时间，比TF-PS节省了52%的时间；在Deep Crossing模型上，PS-Hybrid比PS-Lite节省了33%的训练时间，比TF-PS节省了50%的时间。
4.2.3 混合通信框架扩展性图 8展示了TF-PS、PS-Lite和PS-Hybrid在计算节点数量不断增加的情况下的训练时间变化。实验模型为DCN，数据集为Criteo。从图 8可以看到，在不同计算节点规模下，PS-Lite和PS-Hybrid的训练时间均低于TF-PS。当计算节点个数为4、6时，PS-Hybrid的训练时间与PS-Lite相当，这是因为此时跨网络通信影响了模型训练的整体速度。随着计算节点个数的增加，更多的通信发生在单机内部。当计算节点个数大于8时，PS-Hybrid开始展现出性能优势；当计算节点个数为16时，PS-Hybrid比PS-Lite加速了30%，比TF-PS加速了48%。随着计算节点数量的增加，PS-Hybrid框架的训练时间在大幅下降，这证明混合通信方案具有良好的扩展性。

图 8 不同训练框架的扩展性比较

图选项

4.2.4 稠密参数规模影响当稠密参数规模增大，NCCL的性能优势会更加明显。下面基于Deep Crossing模型和Criteo数据集来验证本文的分析。Deep Crossing模型的稠密参数部分就是多层的残差网络，而残差网络层由2个全连接层组成，因此隐藏层维度调整为原来的2倍，整体的稠密参数规模也调整为原来的2倍，即直接通过调整残差网络的隐藏层维度来调整稠密参数的规模。图 9展示了在隐藏层维度从128调整到1 024过程中，基于PS-Lite和PS-Hybrid框架的Deep Crossing模型在Criteo数据集上的每轮训练时间，残差网络层数为3，计算节点个数为8。

图 9 不同框架的训练时间随稠密参数规模变化

图选项

从图 9可以看出，随着稠密参数量的增大，基于PS-Lite和PS-Hybrid框架的模型训练时间的差距在增大，基于PS-Lite框架的训练时间增加较快，基于PS-Hybrid框架的训练时间增加则不明显。当隐藏层维度为128时，PS-Hybrid较PS-Lite加速33%，当隐藏层维度为1 024时，PS-Hybrid较PS-Lite加速61%。因此，基于混合通信的推荐模型训练框架在稠密参数规模较大时更具优势。
4.2.5 网络通信优化本小节探究3.3节提到的网络通信优化对系统性能的提升，图 10展示了实验结果。本文在Deep Crossing模型上进行实验。图 10的横坐标是梯度的分片粒度，分片粒度从1 000变化到1×10⁷，纵坐标是每轮训练时间。Deep Crossing模型的残差层数设置为7，隐藏层维度设置为256。使用Criteo数据集进行训练，本文从Criteo数据集前500万个样本中采集3万条样本作为训练数据集来测量训练时间。参数服务器个数设置为1，计算节点个数设置为2。

图 10 网络通信优化对系统性能的影响

图选项

本文发现，对于Deep Crossing模型，当梯度的分片粒度为1 000时，模型训练时间为100 s，在分片粒度为1×10⁴和1×10⁵时，模型训练时间分别为95和94 s，当分片粒度为1×10⁷时，模型训练时间升高，每轮训练需要102 s。整体训练时间随着分片粒度的增加同样是先降低后升高。在Deep Crossing模型上，相比FIFO的调度策略，分片增强的优先级通信调度最多加速了7.8%。
5 总结本文通过对推荐模型训练在不同训练架构下的网络通信情况进行定量分析，提出了一个融合参数服务器和全局规约的混合通信分布式训练框架。该框架通过对推荐模型中的稀疏嵌入矩阵和稠密参数采用不同的通信架构，减少了用于参数同步的网络通信开销，很好地利用了GPU的带宽资源。为了进一步优化分布式执行效率，采用优先级传输技术使得计算和通信充分重叠。笔者基于经典的参数服务器PS-Lite和NCCL实现了该混合通信方案，并通过充分的实验验证了所提出的混合通信架构在推荐模型分布式训练场景下，相比于纯参数服务器的方案，在保证模型收敛的前提下能大幅缩短训练时间，具有更好的性能。

参考文献

[1]	FEDUS W, ZOPH B, SHAZEER N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[Z/OL]. (2021-01-11)[2021-05-12]. https://arxiv.org/abs/2101.03961.
[2]	CHENG H T, KOC L, HARMSEN J, et al. Wide & deep learning for recommender systems[C]// Proceedings of the 1st Workshop on Deep Learning for Recommender Systems. Boston, USA, 2016: 7-10.
[3]	WANG R X, FU B, FU G, et al. Deep & cross network for ad click predictions[C]// Proceedings of the ADKDD'17. Halifax, Canada, 2017: 1-7.
[4]	LI M, ANDERSEN D G, PARK J W, et al. Scaling distributed machine learning with the parameter server[C]// Proceedings of the 11th USENIX Symposium on Operating Systems Design and Implementation. Broomfield, USA, 2014: 583-598.
[5]	JIANG J, YU L L, JIANG J W, et al. Angel: A new large-scale machine learning system[J]. National Science Review, 2018, 5(2): 216-236.
[6]	LI M, ANDERSEN D G, SMOLA A, et al. Communication efficient distributed machine learning with the parameter server[C]// Proceedings of the 27th International Conference on Neural Information Processing Systems. Montreal, Canada, 2014, 1: 19-27.
[7]	ABADI M, BARHAM P, CHEN J M, et al. TensorFlow: A system for large-scale machine learning[C]// Proceedings of the 12th USENIX conference on Operating Systems Design and Implementation. Savannah, USA, 2016: 265-283.
[8]	PASZKE A, GROSS S, MASSA F, et al. PyTorch: An imperative style, high-performance deep learning library[C]// Proceedings of the 33rd Conference on Neural Information Processing Systems. Vancouver, Canada, 2019: 8026-8037.
[9]	LI S, ZHAO Y L, VARMA R, et al. PyTorch distributed: Experiences on accelerating data parallel training[J]. Proceedings of the VLDB Endowment, 2020, 13(12): 3005-3018.
[10]	PS-Lite. A light and efficient implementation of the parameter server framework[R/OL]. [2021-05-12]. https://github.com/dmlc/ps-lite.
[11]	NVIDIA Developer. NVIDIA collective communications library (NCCL)[R/OL]. [2021-05-12]. https://developer.nvidia.com/nccl.
[12]	JIANG B Y, DENG C, YI H M, et al. XDL: An industrial deep learning framework for high-dimensional sparse data[C]// Proceedings of the 1st International Workshop on Deep Learning Practice for High-Dimensional Sparse Data. Anchorage, USA, 2019: 1-9.
[13]	PATARASUK P, YUAN X. Bandwidth efficient all-reduce operation on tree topologies[C]// 2007 IEEE International Parallel and Distributed Processing Symposium. Long Beach, USA, 2007: 1-8.
[14]	ANDREW G. Bringing HPC techniques to deep learning [R/OL]. [2021-05-12]. https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/.
[15]	JIA X Y, SONG S T, HE W, et al. Highly scalable deep learning training system with mixed-precision: Training ImageNet in four minutes[C]// Proceedings of the 32nd Conference on Neural Information Processing Systems. Montréal, Canada, 2018.
[16]	Hetu. A high-performance distributed deep learning system targeting large-scale and automated distributed training. [R/OL]. [2021-05-12]. https://github.com/PKU-DAIR/Hetu.
[17]	PENG Y H, ZHU Y B, CHEN Y R, et al. A generic communication scheduler for distributed DNN training acceleration[C]// Proceedings of the 27th ACM Symposium on Operating Systems Principles. Huntsville, Canada, 2019: 16-29.
[18]	SHAN Y, HOENS T R, JIAO J, et al. Deep crossing: Web-scale modeling without manually crafted combinatorial features[C]// Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. San Francisco, USA, 2016: 255-262.