一种晶圆级架构AI加速芯片的任务调度方法及介质

技术领域

本发明涉及云计算技术领域，特别涉及一种晶圆级架构AI加速芯片的任务调度方法及介质。

背景技术

大语言模型是一种基于神经网络的自然语言处理技术，其可以理解、生成人类语言，进而实现智能对话、文本生成、翻译等功能。近年来，大语言模型如生成预训练模型(GPT-3，Generative Pretrained Transformer)等获得了突破性进展，使得及其能够理解人类语言，甚至可以写成一些具有创造性的文本，引领人工智能进入了一个新的阶段。大语言模型基于深度学习技术，利用神经网络对大量文本数据进行训练，在训练过程中，模型会不断地学习语言的规律，包括词汇、语法、以及逻辑等，通过学习逐渐掌握人类语言的复杂性，从而可以生成符合语言规则的文本。GPT-3的核心结构为“Transformer”，是一种基于自注意力机制的特殊神经网络架构，GPT-3的参数量多达1750亿，因此其训练所需要的计算芯片数量可达上千块，且训练时长通常需要数月之久。

为了提高训练效率，有研究提出了一种晶圆级架构。例如特斯拉公司开发的道场(Dojo)超级计算架构，以及Cerebras Systems公司推出的WSE(Wafer Scale Engine)芯片等。通常，所述晶圆级架构中涵盖计算、网络、输入/输出(I/O)芯片到指令集架构(ISA)、电源传输、包装和冷却等功能，可以大规模地运行定制的、特定的机器学习训练算法，实现高速互联的带宽、极低通信延迟、大规模的算力。图1示出一种晶圆级架构AI加速芯片的结构示意图。如图1所示，晶圆级架构AI加速芯片由多个晶粒架构通过芯粒(Chiplet)的先进封装技术封装在晶圆大小的硅基板上。每个晶粒架构中包含若干可相互连接的微架构。如图1所示，所述微架构中例如可包括静态随机存取存储器(Static Random-Access Memory，SRAM)、张量操作模块、张量加速器、在网计算引擎模块、路由模块以及若干RV核等，其可独立或与其他微架构一起执行相关计算任务。

然而目前暂未有公开的关于所述晶圆级架构中任务调度及分配方法，因此，如何提高晶圆级架构的硬件利用率是急需解决的问题。

发明内容

针对现有技术中的部分或全部问题，实现晶圆级架构AI加速芯片的计算任务的自动映射与调度，本发明第一方面提供一种晶圆级架构AI加速芯片的任务调度方法，包括：

将计算任务转换为计算图；

对所述计算图进行切分，以获取至少一个切分方案，每个切分方案包括对所述计算图进行切分得到的至少一个计算子图；

对芯片的计算资源进行划分，以获取至少一个计算资源划分方案，每个计算资源划分方案包括对所述芯片的计算资源进行划分得到的至少一个子网格，且每个子网格包括至少一个计算核；

将不同切分方案和/或计算资源划分方案进行匹配，以获取至少一个匹配方案，每个匹配方案包括一个计算资源划分方案中的各个子网格及其对应的计算子图；以及

根据代价模块，获取一个效率最优的匹配方案。

进一步地，将所述计算图切分包括：

在计算图中插入标记节点，对所述计算图进行切分。

进一步地，对芯片的计算资源进行划分包括：

将所述芯片的计算核划分为至少一个矩形区域，其中每个矩形区域的尺寸相同或不同，且每个矩形区域中包括至少一个计算核。

进一步地，获取一个效率最优的匹配方案包括：

根据代价模型，遍历计算每个匹配方案的代价，并选取代价最小的一个匹配方案执行所述计算任务。

进一步地，获取一个效率最优的匹配方案包括：

训练启发式搜索模型；以及

采用所述启发式搜索模型，基于所述切分方案、计算资源划分方案、匹配方案构成的解空间，获取最优匹配方案。

进一步地，所述任务调度方法还包括：

在所述最优匹配方案的每个子网格内对计算子图进行张量切分，并根据代价模型获取效率最优的任务调度方案。

进一步地，对计算子图进行张量切分包括：

获取所述计算子图中的每个算子的张量信息；以及

根据所述计算子图对应的子网格中计算核的数量，对每个算子各自的张量进行不同维度上的切分，其中每个算子的切分方式相同或不同。

进一步地，所述任务调度方法还包括：

对每个匹配方案中的每个子网格内的计算子图进行张量切分，以获取至少一个任务调度方案；以及

根据代价模型，获取一个效率最优的任务调度方案，其中所述任务调度方案包括计算子图与子网格的匹配方案以及所述匹配方案中各计算子图的张量切分方案。

进一步地，所述代价模型用于评估所述计算任务的执行时间以及内存占用，其中所述执行时间根据如下公式计算得到：

其中，

s

(n

t

B为小批量(micro-batch)的个数。

基于如前所述的计算资源分配方法，本发明第二方面提供一种用于晶圆级架构AI加速芯片的任务调度的计算机可读存储介质，其存储有计算机程序，所述计算机程序在处理器上运行时执行如前所述的任务调度方法。

本发明提供的一种用于晶圆级架构AI加速芯片的任务调度方法，将数据并行、张量并行以及流水线并行等多种并行方式有机结合起来，分别切分计算任务、计算资源，最终将其进行匹配，并可进一步对计算子图进行张量切分。其每个步骤均具有寻优的空间，考虑的优化空间更大更全面，从而能够有效地提高晶圆级架构AI加速芯片的硬件利用效率。同时所述任务调度方法支持不同优化算法的拓展，相较于手动方法而言，其无需专家经验，适应性强，具有较强的通用性和扩展性，可广泛应用于不同类型的计算任务和计算资源，且可拓展至其他一般的计算架构，比如GPU集群、TPU等专用加速器组成的集群等。

附图说明

为进一步阐明本发明的各实施例的以上和其它优点和特征，将参考附图来呈现本发明的各实施例的更具体的描述。可以理解，这些附图只描绘本发明的典型实施例，因此将不被认为是对其范围的限制。在附图中，为了清楚明了，相同或相应的部件将用相同或类似的标记表示。

图1示出一种晶圆级架构AI加速芯片的结构示意图；

图2a-2c分别示出现有的三种并行方式的过程示意图；

图3示出本发明一个实施例的解空间的示意图；

图4示出本发明一个实施例的用于晶圆级架构AI加速芯片的任务调度方法的过程示意图；

图5示出本发明一个实施例的用于晶圆级架构AI加速芯片的任务调度方法的流程示意图；

图6示出本发明一个实施例的计算图的切分方法示意图；

图7示出本发明一个实施例的计算资源的划分方法示意图；以及

图8示出本发明一个实施例的张量切分方法示意图。

具体实施方式

以下的描述中，参考各实施例对本发明进行描述。然而，本领域的技术人员将认识到可在没有一个或多个特定细节的情况下或者与其它替换和/或附加方法或组件一起实施各实施例。在其它情形中，未示出或未详细描述公知的结构或操作以免模糊本发明的发明点。类似地，为了解释的目的，阐述了特定数量和配置，以便提供对本发明的实施例的全面理解。然而，本发明并不限于这些特定细节。此外，应理解附图中示出的各实施例是说明性表示且不一定按正确比例绘制。

在本说明书中，对“一个实施例”或“该实施例”的引用意味着结合该实施例描述的特定特征、结构或特性被包括在本发明的至少一个实施例中。在本说明书各处中出现的短语“在一个实施例中”并不一定全部指代同一实施例。

需要说明的是，本发明的实施例以特定顺序对方法步骤进行描述，然而这只是为了阐述该具体实施例，而不是限定各步骤的先后顺序。相反，在本发明的不同实施例中，可根据实际需求的调节来调整各步骤的先后顺序。

为了提高晶圆级架构芯片的硬件利用率，本发明将晶圆级架构的AI任务自动调度与映射抽象成一个可解的优化问题，即训练一个模型从定义的解空间中获取最优解。一般的模型训练通常使用数据并行，大模型训练一般还会使用模型并行，其中模型并行包括张量并行、流水线并行两种，图2a-2c分别示出现有的三种并行方式的过程示意图，如图所示，数据并行是将数据进行拆分同时输入多个设备并行计算，张量并行则是将张量操作划分到多个设备，以加速计算或增加模型大小，而流水线并行则是在张量并行的基础上进一步引入数据并行的办法，将原先的数据划分成若干个小批量(micro-batch)后再逐一送入芯片、设备进行训练。而为了有效混合使用多种并行方式，进而提高效率，在本发明中将几种并行方式有机结合起来，设计了4个步骤：计算图切分成计算子图、计算资源切分成子网格、计算子图与子网格的匹配、子网格内对子图的张量切分，每个步骤都具有寻优的空间，且这四个步骤与本发明定义的解空间相对应。图3示出本发明一个实施例的解空间的示意图。如图3所示，在本发明的实施例中，所述解空间包括计算图的切分方案、计算资源的划分方案、切分得到的计算子图与划分得到子网格之间的匹配方案、以及子网格内对计算子图的张量切分方案。

基于此，图4示出本发明一个实施例的用于晶圆级架构AI加速芯片的任务调度方法的过程示意图。如图4所示，所述任务调度方案以计算任务及晶圆级架构的计算资源作为输入，通过优化求解算法在如前所述的解空间中需求最优映射，进而通过路由、调度、内存分配与管理最终完成计算任务，在需求最优解的过程中，可通过代价模型进行代价分析。

为进一步地说明所述任务调度方法的具体实施过程，图5示出本发明一个实施例的用于晶圆级架构AI加速芯片的任务调度方法的流程示意图，如图5所示，一种用于晶圆级架构AI加速芯片的任务调度方法，包括：

首先，在步骤501，获取计算图。将计算任务转换为计算图，计算图能描述接收待处理的计算任务，进而分析计算任务的特征，包括计算任务的计算类型、计算量、数据依赖关系、数据形状等。在本发明的实施例中，可以采用主流AI框架如Tensorflow等生成计算图；

接下来，在步骤502，切分计算图。对所述计算图进行切分，以获取至少一个切分方案，每个切分方案包括对所述计算图进行切分得到的至少一个计算子图。在本发明的一个实施例中，如图6所示，所述计算图的切分通过在AI框架自有的计算图中插入标记节点，然后通过编写对应的切分变换程序实现切分，进而得到若干个计算子图；

同时，在步骤503，划分计算资源。对芯片的计算资源进行划分，以获取至少一个计算资源划分方案，每个计算资源划分方案包括对所述芯片的计算资源进行划分得到的至少一个子网格，且每个子网格包括至少一个计算核。如图1所示，一个晶圆级架构AI加速芯片的计算资源可以抽象成m×n个计算核(tile)，所述m×n个计算核构成一个设备，则如图7所示，计算资源的划分就是对所述设备进行划分，将其分为不同矩形区域，其中每个矩形区域即为一个子网格，每个子网格包括若干计算核，每个矩形区域的尺寸可以详谈或不同，即每个子网格中包含的计算核数量可以相同或不同；

接下来，在步骤504，匹配子网格与计算子图。将不同切分方案和/或计算资源划分方案进行匹配，以获取至少一个匹配方案，子网格与计算子图的匹配意味着每个子网格区域各自执行自己分配到的计算子图任务，即每个匹配方案包括一个计算资源划分方案中的各个子网格及其对应的计算子图；

此外，为了更好地提高硬件利用率，在本发明的一个实施例中，在匹配方案的基础上还可以进一步地进行步骤505，张量切分。在子网格对计算子图进行张量切分，即描述所述计算子图在子网格内具体是怎么执行的。如图8所示，子网格内计算子图的张量切分通过设置计算子图中每个算子的切分方式实现。例如可以通过一些标记语言来描述参与计算的数据是否进行切分。如图8所示，对于一个三维的张量形状为[b,h,w]，当其与形状为[m0,m1]的子网格匹配时，则每个算子可以采用以下三种切分方式中的一种：沿着计算网格的m0维度切分，采用S0表示；沿着计算网格的m1维度切分，S1表示；以及用R表示不切分，则所述三维张量的切分可以被描述为[b维度的切分，h维度的切分，w维度的切分]，例如[S1,R,S0]就是可能的一种切分方案；以及

最后，在步骤506，优化求解。如前所述，在前述步骤中获取的切分方案、计算资源划分方案以及匹配方案、和/或张量切分方案可构成解空间，最终需要从解空间中寻求最优解，实现对抽象出来的优化问题的求解。在本发明的实施例中，最终优化求解可以使用遍历搜索算法，也可以使用启发式搜索算法比如遗传算法。其中遍历搜索算法是指根据代价模型对每个匹配方案或张量切分方案进行评价，并选取代价最小的一个座位最终的任务调度方案。而启发式搜索算法是指通过训练模型，并采用模型在解空间中寻求最优解。

在本发明的一个实施例中，所述解空间可仅由切分方案、计算资源划分方案以及匹配方案构成，在这一实施例中，首先需要通过如前所述的优化算法确定最优的匹配方案，然后在所述最优匹配方案的每个子网格内对计算子图进行张量切分，并根据代价模型获取效率最优的任务调度方案。

在本发明的另一个实施例中，所述解空间则包括切分方案、计算资源划分方案、匹配方案、以及张量切分方案，直接通过如前所述的优化算法确定最终的任务调度方案，在这一实施例中，任务调度方案包括计算子图与子网格的匹配方案以及所述匹配方案中各计算子图的张量切分方案。

在本发明的一个实施例中，为了更准确地对方案进行评价，还设置有代价模型。所述代价模型是多层次的，不同层次包括单个计算子图在子计算资源上的执行、多个子图的流水线并行执行、多晶圆级硬件集群上多子图的流水线执行等。因此所述代价模型一方面可以进行计算、通信、访存的建模，另一方面还可精准高效地反馈执行时间、内存占用等信息。具体而言，在本发明的一个实施例中，所述代价模型是使用真实硬件评估反馈。在本发明的又一个实施例中，所述代价模型则是收集真实硬件的反馈构建数据集，训练一个模型来学习、拟合反馈。在本发明的再一个实施例中，所述代价模型则是通过构建仿真器来建模仿真获得等等。建模仿真的方法需要能够评估最底层的计算和通信，继而能够评估每个计算和通信算子的执行时间与内存占用，进而能够计算出每个计算子图在子网格上的耗时与内存，进而得到整个计算任务的整体耗时：

其中，

s

(n

t

B为小批量(micro-batch)的个数。

至此可完成晶圆级架构AI加速芯片的任务调度，在运行时，根据得到的任务调度方案对硬件资源进行调度，例如包括路由功能的支持、计算通信功能的调度、内存管理功能等。

所述任务调度方法可是实现自动化的AI任务在晶圆级架构上的映射与调度，相较于手动方法一般需要专家手动调整，进而很长时间才能针对某一特定AI任务得到比较好的映射而言，所述任务调度方法更为快速、适应性更强，且不需要专家经验。此外，所述任务调度方案考虑的优化空间更大更全面，从而能够得到的映射方案，硬件利用效率更高，同时支持不同优化算法的拓展。同时，所述任务调度方法具有较强的通用性和扩展性，可广泛应用于不同类型的计算任务和计算资源。

基于如前所述的计算资源分配方法，本发明还提供一种用于晶圆级架构AI加速芯片的任务调度的计算机可读存储介质，其存储有计算机程序，所述计算机程序在处理器上运行时执行如前所述的任务调度方法。

尽管上文描述了本发明的各实施例，但是，应该理解，它们只是作为示例来呈现的，而不作为限制。对于相关领域的技术人员显而易见的是，可以对其做出各种组合、变型和改变而不背离本发明的精神和范围。因此，此处所公开的本发明的宽度和范围不应被上述所公开的示例性实施例所限制，而应当仅根据所附权利要求书及其等同替换来定义。