面向分布式深度学习算子并行训练的通信优化方法与系统

技术领域

本发明涉及分布式深度学习并行训练技术领域，尤其涉及一种面向分布式深度学习算子并行训练的通信优化方法与系统。

背景技术

2021年美国斯坦福大学，微软公司和英伟达公司的深度学习团队在SC（International Conference for High Performance Computing, Networking, Storageand Analysis，高性能计算、网络、存储和分析国际会议）会议上发表用算子并行的方法训练大型语言模型的系统Megatron，并且在开源社区GitHub上开源。为了训练单张GPU卡的内存无法容纳的大模型，以及加快模型训练的速度，该方法把一个计算算子拆分到多块GPU卡上并行执行，接着串行地执行一个通信算子，同步各个GPU的计算结果，以保证和原来单计算算子的计算逻辑一致性。但是该方法中，在一台服务器（包括四张NVIDIA1080Ti GPU）的环境下做算子并行训练，通信开销是在影响训练速度的关键路径上。本质原因在于以下三点：1）通信算子必须等待各个卡的计算完成后，才能开始执行，并且集合通信执行完成后才能继续开始后续的计算，这导致引入的通信开销无法被隐藏；2）模型训练无法充分利用系统中网络和计算资源，计算算子执行时，网络处于空闲状态；3）为了减小引入的通信开销，Megatron系统优化了算子切分策略，但是在一个变压器（Transformer）中仍然需要两次集合通信操作，且每次的通信量和模型的大小成正比。在训练420M-BERT模型时，Transformer层每训练一个数据需要的通信量为1GB，通信时间约占端到端训练时间的30%。

发明内容

本发明的目的是提供一种面向分布式深度学习算子并行训练的通信优化方法与系统，对通信的优化方法，以避免现有技术的上述缺点，在不影响训练精度的前提下，将通信开销隐藏到计算中，提升系统中网络和计算资源的利用率。

本发明的目的是通过以下技术方案实现的：

一种面向分布式深度学习算子并行训练的通信优化方法，包括：

获取模型训练框架的初始计算流图，并确定模型算子并行训练中各计算算子的计算量与计算完成后的对应通信算子的通信量；

预先构造不同计算量的计算算子，以及不同通信量的通信算子，并分别执行，获得计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式，结合计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式对各计算算子的计算量与计算完成后的对应通信算子的通信量进行不同拆分粒度的切分，选出最优拆分粒度；

对于模型算子并行训练中每一计算算子及其计算量，以及每一通信算子及其通信量，均按照所述最优拆分粒度进行拆分，拆分结果为与所述最优拆分粒度相同数目的子计算算子、每一子计算算子对应的计算量、与所述最优拆分粒度相同数目的子通信算子，以及每一通信算子对应的通信量，其中子计算算子与子通信算子并行执行；

按照拆分结果对所述初始计算流图进行修改，由所述模型训练框架按照修改后的计算流图执行模型训练。

一种面向分布式深度学习算子并行训练的通信优化系统，包括：

剖析器，用于获取模型训练框架的初始计算流图，并确定模型算子并行训练中各计算算子的计算量与计算完成后的对应通信算子的通信量，计算算子与通信算子一一对应；

决策者，用于预先构造不同计算量的计算算子，以及不同通信量的通信算子，并分别执行，获得计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式，结合计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式对各计算算子的计算量与计算完成后的对应通信算子的通信量进行不同拆分粒度的切分，选出最优拆分粒度；

重叠器，用于对于模型算子并行训练中每一计算算子及其计算量，以及每一通信算子及其通信量，均按照所述最优拆分粒度进行拆分，拆分结果为与所述最优拆分粒度相同数目的子计算算子、每一子计算算子对应的计算量、与所述最优拆分粒度相同数目的子通信算子，以及每一通信算子对应的通信量，其中子计算算子与子通信算子并行执行；

改写器，用于按照拆分结果对所述初始计算流图进行修改，由所述模型训练框架按照修改后的计算流图执行模型训练。

一种处理设备，包括：一个或多个处理器；存储器，用于存储一个或多个程序；

其中，当所述一个或多个程序被所述一个或多个处理器执行时，使得所述一个或多个处理器实现前述的方法。

一种可读存储介质，存储有计算机程序，当计算机程序被处理器执行时实现前述的方法。

由上述本发明提供的技术方案可以看出，对通信方式进行了优化，传统方法相比，利用开销模型评估出最优的拆分策略，将通信前后的计算算子拆分成细粒度的子计算算子，使得计算算子和通信算子并行，使得通信开销能最大化地隐藏到计算中，最大化利用网络和计算资源。

附图说明

为了更清楚地说明本发明实施例的技术方案，下面将对实施例描述中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附图仅仅是本发明的一些实施例，对于本领域的普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据这些附图获得其他附图。

图1为本发明实施例提供的一种面向分布式深度学习算子并行训练的通信优化方法的流程图；

图2为本发明实施例提供的算子并行实现矩阵乘法的样例示意图；

图3为本发明实施例提供的不同拆分粒度下的算子执行开销样例示意图；

图4为本发明实施例提供的子计算算子与子通信算子并行执行的原理图；

图5为本发明实施例提供的一种面向分布式深度学习算子并行训练的通信优化方法架构以及运行流程示意图；

图6为本发明实施例提供的一种面向分布式深度学习算子并行训练的通信优化系统的示意图；

图7为本发明实施例提供的一种处理设备的示意图。

具体实施方式

下面结合本发明实施例中的附图，对本发明实施例中的技术方案进行清楚、完整地描述，显然，所描述的实施例仅仅是本发明一部分实施例，而不是全部的实施例。基于本发明的实施例，本领域普通技术人员在没有做出创造性劳动前提下所获得的所有其他实施例，都属于本发明的保护范围。

首先对本文中可能使用的术语进行如下说明：

术语“和/或”是表示两者任一或两者同时均可实现，例如，X和/或Y表示既包括“X”或“Y”的情况也包括“X和Y”的三种情况。

术语“包括”、“包含”、“含有”、“具有”或其它类似语义的描述，应被解释为非排它性的包括。例如：包括某技术特征要素（如原料、组分、成分、载体、剂型、材料、尺寸、零件、部件、机构、装置、步骤、工序、方法、反应条件、加工条件、参数、算法、信号、数据、产品或制品等），应被解释为不仅包括明确列出的某技术特征要素，还可以包括未明确列出的本领域公知的其它技术特征要素。

下面对本发明所提供的一种面向分布式深度学习算子并行训练的通信优化方案进行详细描述。本发明实施例中未作详细描述的内容属于本领域专业技术人员公知的现有技术。本发明实施例中未注明具体条件者，按照本领域常规条件或制造商建议的条件进行。

实施例一

本发明实施例提供一种面向分布式深度学习算子并行训练的通信优化方法，如图1所示，主要包括：

步骤1、剖析模型。

本步骤中，通过获取模型训练框架的初始计算流图，从而确定模型算子并行训练中各计算算子的计算量与计算完成后的对应通信算子的通信量。

本发明实施例中，通过定位需要通信点（通信位置），确定通信点前计算算子的计算量，以及计算完成后的通信量。其中，计算量可以使用执行完相关计算算子所需要的浮点数运算次数表示。

如图2所示，一个矩阵乘MatMul算子（计算算子）被切分到

图2中：X表示输入数据，Y

步骤2、策略评估和生成最佳优化策略。

本发明实施例中，预先构造不同计算量的计算算子，以及不同通信量的通信算子，并分别执行，获得计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式，结合计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式对各计算算子的计算量与计算完成后的对应通信算子的通信量进行不同拆分粒度的切分，选出最优拆分粒度。具体来说：

构造不同计算量的计算算子，以及不同通信量的通信算子，遍历每个算子，对于每个算子，在其执行前后加上时间戳，执行若干次后，分别统计每一次执行需要的时间并记录，再求取平均值，作为执行相应算子所需时间，所述每个算子包括：计算算子与通信算子，执行计算算子所需时间即为执行对应计算量所需时间，执行通信算子所需时间即为执行对应通信量所需时间；执行完所有不同计算量的计算算子后，获得执行不同计算量所需时间，执行完所有不同通信量的通信算子后，获得执行不同通信量所需时间，再通过数据拟合，获得计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式。

通过尝试不同拆分粒度，并通过评估函数（开销模型）对不同拆分粒度进行评估，评估函数表示为：

其中，Cost表示评估值，C表示计算量，N表示通信量；S表示拆分粒度，即拆分后的份数；

本发明实施例中，计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式可以预先确定，在实际应用中，结合上述评估函数，为每一对计算算子的计算量与计算完成后的对应通信算子的通信量分别选出最优拆分粒度。

示例性的：在四张NVIDIA 1080Ti GPU上对维度为（512，8，2048）*（2048，2048）=（512，8，2048）的矩阵乘做算子拆分粒度的评估，结果如图3所示，表示不同的拆分粒度下执行完计算算子的逻辑需要的开销，其中，将单张GPU卡上做算子并行的算子拆分成16份时Cost值最小，即性能最优。

步骤3、算子拆分和隐藏通信。

本发明实施例中，对于模型算子并行训练中每一计算算子及其计算量，以及每一通信算子及其通信量，均按照所述最优拆分粒度进行拆分，拆分结果为与所述最优拆分粒度相同数目的子计算算子、每一子计算算子对应的计算量、与所述最优拆分粒度相同数目的子通信算子，以及每一通信算子对应的通信量，其中子计算算子与子通信算子并行执行。此处的原理为：确定最优拆分粒度（记为S*）进行拆分后，将计算算子的第一维拆分成S*个子计算算子，在每个子计算算子计算完成后立即做集合通信，因此通信算子也被拆分为S*个子通信算子，由于子计算算子和其他子计算算子的通信之间没有逻辑依赖，所以可以使子计算算子和其他子计算算子对应的通信并行执行。

示例性的：假设最优拆分粒度S*=16，需要将计算（512，8，2048）*（2048，2048）按照第一维拆分成16个子计算算子，每一子计算算子为（32，8，2048）*（2048，2048），由于每个计算算子计算完成后都需要进行一次通信，拆分后的子计算算子也是如此，因此，计算算子被拆分为16个子计算算子后，通信算子自然也被拆分为子通信算子，未拆分前计算（512，8，2048）*（2048，2048）=（512，8，2048），计算结果（512，8，2048）的大小=512*8*2048*4/1024/1024=32MB，其中4表示一个数字占的字节数，也就是说，未拆分前通信量为一次32MB，计算算子被拆分为16个子计算算子后，通信算子自然也被拆分为16个子通信算子，也就是每次通信量变为2MB。

实现计算和通信并行优选方式为：分配两个任务流（CUDA Stream），称为计算流（负责计算）与通信流（负责通信），将每个子计算算子依次交给计算流执行，每个子计算算子执行完成后，将对应的子通信算子交给通信流执行，并且在计算流中开始下一个子计算算子的执行。

如图4所示，提供了子计算算子与子通信算子并行执行的原理图，图4中以最优拆分粒度S*=2为例，其中，原子算子中计算

原始算子执行流程如下：计算

算子拆分后，子计算算子与子通信算子并行执行流程如下：执行

步骤4、模型训练与算子拆分的集成。

按照拆分结果对所述初始计算流图进行修改，由所述模型训练框架按照修改后的计算流图执行模型训练。具体的，根据前文获得的拆分结果，把模型在前向计算和反向计算中做算子并行的算子替换为上述拆分后的子计算组和子通信组，把修改后计算流图交给模型训练框架执行，模型训练框架执行的过程可参照常规技术，故不做赘述。

如图5所示，展示了本发明的整体架构以及运行流程示意图。其中，Profiler（剖析器）用于执行前述步骤1，Planner（决策者）用于执行前述步骤2，Overlap（重叠器）用于执行前述步骤3，Rewrite（改写器）用于执行前述步骤4；①表示最初阶段从PyTorch Dataflow（PyTorch数据流）中获取初始计算流图，②表示将各计算算子的计算量与计算完成后的对应通信算子的通信量传输至Planner，③表示将最优拆分粒度传输至Overlap，④表示将拆分结果传输至Rewrite，⑤表示将修改后的计算流图传输至PyTorch Engine（PyTorch框架引擎）。

本发明实施例提供的上述方案，在分布式深度学习算子并行训练场景下，对通信方式进行了优化，与传统方法相比，利用开销模型评估出最优的拆分策略，将通信前后的计算算子拆分成细粒度的子算子，使得计算算子和通信算子并行，使得通信开销能最大化地隐藏到计算中，提高了模型训练的速度，最大化利用网络和计算资源。

通过以上的实施方式的描述，本领域的技术人员可以清楚地了解到上述实施例可以通过软件实现，也可以借助软件加必要的通用硬件平台的方式来实现。基于这样的理解，上述实施例的技术方案可以以软件产品的形式体现出来，该软件产品可以存储在一个非易失性存储介质（可以是CD-ROM，U盘，移动硬盘等）中，包括若干指令用以使得一台计算机设备（可以是个人计算机，服务器，或者网络设备等）执行本发明各个实施例所述的方法。

实施例二

本发明还提供一种面向分布式深度学习算子并行训练的通信优化系统，其主要基于前述实施例提供的方法实现，如图6所示，该系统主要包括：

剖析器，用于获取模型训练框架的初始计算流图，并确定模型算子并行训练中各计算算子的计算量与计算完成后的对应通信算子的通信量，计算算子与通信算子一一对应；

决策者，用于预先构造不同计算量的计算算子，以及不同通信量的通信算子，并分别执行，获得计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式，结合计算量与计算时间的函数关系式，以及通信量与通信时间的函数关系式对各计算算子的计算量与计算完成后的对应通信算子的通信量进行不同拆分粒度的切分，选出最优拆分粒度；

重叠器，用于对于模型算子并行训练中每一计算算子及其计算量，以及每一通信算子及其通信量，均按照所述最优拆分粒度进行拆分，拆分结果为与所述最优拆分粒度相同数目的子计算算子、每一子计算算子对应的计算量、与所述最优拆分粒度相同数目的子通信算子，以及每一通信算子对应的通信量，其中子计算算子与子通信算子并行执行；

改写器，用于按照拆分结果对所述初始计算流图进行修改，由所述模型训练框架按照修改后的计算流图执行模型训练。

本发明实施例中，所述构造不同计算量的计算算子，以及不同通信量的通信算子，并分别执行，获得计算量与计算时间，以及通信量与通信时间的函数关系式包括：遍历每个算子，对于每个算子，在其执行前后加上时间戳，执行若干次后，统计执行需要的时间并记录，所述每个算子包括：计算算子与通信算子；执行完所有不同计算量的计算算子后，获得执行不同计算量所需时间，执行完所有不同通信量的通信算子后，获得执行不同通信量所需时间，再通过数据拟合，获得计算量与计算时间，以及通信量与通信时间的函数关系式。

本发明实施例中，所述对各计算算子与通信算子进行不同拆分粒度的切分，选出最优拆分粒度包括：

通过评估函数对不同拆分粒度进行评估，评估函数表示为：

其中，Cost表示评估值，C表示计算量，N表示通信量；S表示拆分粒度，即拆分后的份数；

选出使评估值Cost最小的拆分粒度作为最优拆分粒度。

本发明实施例中，所述子计算算子与子通信算子并行执行包括：分配两个任务流，称为计算流与通信流，将每个子计算算子依次交给计算流执行，每个子计算算子执行完成后，将对应的子通信算子交给通信流执行，并且在计算流中开始下一个子计算算子的执行。

所属领域的技术人员可以清楚地了解到，为描述的方便和简洁，仅以上述各功能模块的划分进行举例说明，实际应用中，可以根据需要而将上述功能分配由不同的功能模块完成，即将系统的内部结构划分成不同的功能模块，以完成以上描述的全部或者部分功能。

实施例三

本发明还提供一种处理设备，如图7所示，其主要包括：一个或多个处理器；存储器，用于存储一个或多个程序；其中，当所述一个或多个程序被所述一个或多个处理器执行时，使得所述一个或多个处理器实现前述实施例提供的方法。

进一步的，所述处理设备还包括至少一个输入设备与至少一个输出设备；在所述处理设备中，处理器、存储器、输入设备、输出设备之间通过总线连接。

本发明实施例中，所述存储器、输入设备与输出设备的具体类型不做限定；例如：

输入设备可以为触摸屏、图像采集设备、物理按键或者鼠标等；

输出设备可以为显示终端；

存储器可以为随机存取存储器（Random Access Memory，RAM），也可为非不稳定的存储器（non-volatile memory），例如磁盘存储器。

实施例四

本发明还提供一种可读存储介质，存储有计算机程序，当计算机程序被处理器执行时实现前述实施例提供的方法。

本发明实施例中可读存储介质作为计算机可读存储介质，可以设置于前述处理设备中，例如，作为处理设备中的存储器。此外，所述可读存储介质也可以是U盘、移动硬盘、只读存储器（Read-Only Memory，ROM）、磁碟或者光盘等各种可以存储程序代码的介质。

以上所述，仅为本发明较佳的具体实施方式，但本发明的保护范围并不局限于此，任何熟悉本技术领域的技术人员在本发明披露的技术范围内，可轻易想到的变化或替换，都应涵盖在本发明的保护范围之内。因此，本发明的保护范围应该以权利要求书的保护范围为准。