1.1.1 基本概念
分布式机器学习是指部署在多个计算机节点,根据特定的通信范式和同步协议,通过通信网络交换机器学习模型数据,实现多机数据协同和分布式并行计算的一类大规模机器学习训练方法。
分布式机器学习系统整体架构概览如图1-2所示。分布式机器学习系统包含三个核心部件:数据资源池、计算节点群和通信网络。数据资源池将大数据分散存储在多台计算机的文件系统以实现大存储,可为计算节点群提供单机可承载的数据切片以供训练。计算节点群利用分布式技术汇聚多台计算机的计算资源,并行加速大规模计算任务,为大数据计算提供大算力支持。另外,计算节点群也可以将大模型切片,使得单台计算机有能力处理这些模型切片。最后,多台计算机之间借由通信网络交换计算结果,实现模型参数的同步或流水线式的大模型计算。
图1-2 分布式机器学习系统整体架构概览
在早期尝试中,分布式机器学习主要用在一台计算机上实现多卡协作,这里的卡通常指GPU计算卡,如图1-3(a)所示。一方面,为了整合多卡显存承载大模型,大模型将被切成多个较小的分片,分别缓存在多卡的显存中,并借助CPU调度处理顺序以及传递计算结果。另一方面,多卡并行处理不同批次的数据,再由CPU同步计算结果,可以整合多卡算力实现训练加速。由于CPU和GPU位于同一台物理机上,二者通过高速串行计算机扩展总线(Peripheral Component Interconnect express,PCIe)交换数据,内存与显存之间传输带宽可达64GB/s。在使用NVLink桥接器后,GPU之间、GPU与CPU之间的传输带宽更是可以达到50~200GB/s[22]。众多智能云平台都提供了配备多卡的GPU云服务器,如Azure和腾讯云都提供了1~8卡的配置选项,个体研究者可以租用云算力来满足训练需求。然而,受到机箱容量、主板可用PCIe插槽数量以及电源供电功率等制约,单台机器可以容纳的GPU数量非常有限,通常仅有1~8卡。即便使用8个Tesla V100 GPU,训练GPT-3模型仍需要36年[23]。因此,单机多卡模式可扩展性差,加速性能有限,无法解决数据量大的问题。
图1-3 单机多卡与多机多卡的配置概览
随着人工智能的飞速发展,训练数据量日益庞大,训练大模型的时间也越来越长,单机多卡协作已经无法满足业界的训练需求,这迫使业界探索多机协作模式,使用总计成百上千个GPU的大规模计算集群进一步加速机器学习训练。为简化说明,图1-3(b)所示为仅在服务器内画出一个GPU,但实际也可以像图1-3(a)所示配置1~8卡。在多机协作模式下,为实现GPU之间通信,CPU需要先将GPU显存中的模型数据通过PCIe复制到内存,再经由数据中心网络发送给另一台计算机,另一台计算机将接收到的模型数据复制到GPU显存。上述模型传输过程涉及两个延迟:基于PCIe的GPU显存到CPU内存的数据搬移延迟,以及基于数据中心网络的主机间数据传输延迟。为了消除GPU和CPU之间的数据搬移延迟,可以借助GPU远程直接内存访问(GPUDirect Remote Direct Memory Access,GPUDirect RDMA)[24]技术,该技术将GPU显存数据读写完全交由GPUDirect RDMA硬件执行,数据交换无须经过CPU内存,主机A的GPU1可以直接访问主机B的GPU2显存,从而实现高带宽、低延迟的GPU通信。但是,在高性能计算应用中使用GPUDirect RDMA技术依赖于InfiniBand网络,所需的专用硬件设备成本高昂,一般数据中心难以承受。多数情况下,网络交换机的传输带宽远低于PCIe带宽,使得多机通信成为分布式机器学习的主要性能瓶颈。本书也将围绕多机通信瓶颈进行展开,并介绍一系列通信优化技术。