用于人工智能加速器集群的集合通信方法和装置与流程

本公开涉及人工智能领域，尤其涉及人工智能加速器、数据集合通信等领域，特别是涉及一种用于人工智能加速器集群的集合通信方法、装置、电子设备、计算机可读存储介质和计算机程序产品。

背景技术：

1、随着大语言模型（llm）和深度学习技术的快速发展，人工智能（ai）模型参数量呈现指数级增长趋势，从数十亿增长到万亿规模，单张ai加速卡的显存与算力已无法满足训练或推理需求，必须采用单机多卡乃至多机多卡协同计算。在多卡协同计算的场景中，ai加速器之间的数据通信成为影响整体性能的关键瓶颈。

技术实现思路

1、本公开提供了一种用于人工智能加速器集群的集合通信方法、装置、电子设备、计算机可读存储介质和计算机程序产品。

2、根据本公开的一方面，提供了一种用于人工智能加速器集群的集合通信方法，人工智能加速器集群中的每张加速卡包括多个计算单元，该方法包括：从每张加速卡中的多个计算单元中选取一个计算单元作为主计算单元，除主计算单元之外的至少一个计算单元作为从计算单元；基于所有加速卡的主计算单元构建单向逻辑环作为环维度，基于同一加速卡上的所有计算单元构建包括父子关系的组作为组维度；根据每个主计算单元在环维度上的位置以及该主计算单元是否具有跨卡或跨机网络通信能力，对环维度上的各个主计算单元进行分类；生成集合通信操作的通信执行序列，通信执行序列包括多个按时间顺序排列的硬件可执行的操作指令，操作指令至少包括直接存储器存取dma传输指令、同步事件轮询指令、同步事件提交指令和专用计算引擎归约指令；以及基于通信执行序列中的操作指令构建为多个并行流水线，并部署到人工智能加速器集群的硬件上执行，其中，对于每个通信轮次，通信执行序列指示每个主计算单元在环维度上与该主计算单元的前驱主计算单元和后继主计算单元执行跨卡数据收发，同时在同一通信轮次内，通信执行序列指示同一加速卡内的计算单元在组维度上沿父子关系执行片内数据归约和复制，使得环维度的跨卡通信和组维度的片内通信在时间上至少部分重叠地并行执行。

3、在一些实施例中，基于所有加速卡的主计算单元构建单向逻辑环作为环维度包括：通过各计算单元之间的拓扑信息交换获取每个计算单元的硬件属性，硬件属性包括计算单元标识、所在卡标识、所在节点标识和设备索引；通过搜索算法在候选环排列中搜索使得相邻主计算单元之间通信带宽最大化的最优环排列；以及根据最优环排列确定环维度上每个主计算单元的前驱主计算单元索引和后继主计算单元索引，从而构建单向逻辑环作为环维度。

4、在一些实施例中，根据每个主计算单元在环维度上的位置以及该计算单元是否具有跨卡或跨机网络通信能力，对环维度上的各个主计算单元进行分类包括：当主计算单元位于环维度的首位且具有从网络接收数据的能力而不具有向网络发送数据的能力时，将该计算单元分类为环头部计算单元；当主计算单元位于环维度的末尾且具有向网络发送数据的能力而不具有从网络接收数据的能力时，将该计算单元分类为环尾部计算单元；当主计算单元既不从网络接收数据也不向网络发送数据时，将该计算单元分类为环身体计算单元；以及当主计算单元同时具有从网络接收数据和向网络发送数据的能力时，将该计算单元分类为环混合计算单元。

5、在一些实施例中，从每张加速卡中的多个计算单元中选取一个计算单元作为主计算单元，除主计算单元之外的至少一个计算单元作为从计算单元包括：主计算单元为计算单元标识为0的计算单元，从计算单元为计算单元标识不为0的计算单元。

6、在一些实施例中，基于同一加速卡上的所有计算单元构建包括父子关系的组作为组维度包括：遍历同一加速卡上的所有计算单元并按计算单元标识顺序排列；以及在计算单元标识相邻的计算单元之间建立父子关系，其中计算单元标识较小的计算单元作为父计算单元，计算单元标识较大的计算单元作为子计算单元，每个计算单元记录该计算单元的父计算单元索引和子计算单元索引。

7、在一些实施例中，集合通信操作为归约操作，通信执行序列包括：环维度中的每个主计算单元将本地发送缓冲区中的数据沿环维度的方向发送给各自的后继主计算单元；在环维度上执行多轮数据归约转发，其中，每轮中每个主计算单元将从该主计算单元的前驱主计算单元接收的数据块与本地的对应数据块经由专用计算引擎进行归约后发送给该主计算单元的后继主计算单元，同时，在每个通信轮次中，组维度上每个从计算单元与该从计算单元的父计算单元执行片内归约操作；以及各主计算单元将归约结果数据块沿环维度广播给所有其他主计算单元，并在组维度上将归约结果数据块复制到各从计算单元。

8、在一些实施例中，在环维度上执行多轮数据归约转发时，每个主计算单元执行以下步骤：轮询前驱主计算单元的同步事件以确认待归约数据已就绪；启动专用计算引擎对本地数据块和待归约数据块执行归约操作，并将归约结果写入目标缓冲区；对于非最后两轮通信轮次，将目标缓冲区中的归约结果通过dma传输发送给该主计算单元的后继主计算单元，并通过提交同步事件通知该主计算单元的后继主计算单元；以及对于最后两轮通信轮次，仅执行归约操作而不向后继主计算单元转发。

9、在一些实施例中，当人工智能加速器集群包括多个节点时，方法还包括：根据环维度的最优环排列确定环维度的环头部计算单元和环尾部计算单元；为环头部计算单元选择与该环头部计算单元所在节点相连的、带宽最高的网络设备作为数据接收端，并为环尾部计算单元选择与该环尾部计算单元所在节点相连的、带宽最高的网络设备作为数据发送端；以及当环头部计算单元和环尾部计算单元位于同一节点、共用同一网络设备并且环头部计算单元和环尾部计算单元之间的路径跨越多个pcie交换机时，为环头部计算单元和环尾部计算单元选择不同的网络设备以进行跨机网络通信。

10、在一些实施例中，基于通信执行序列中的操作指令构建为多个并行流水线包括：将通信执行序列中的连续dma传输操作构建为拷贝流水线；将环维度上的跨卡通信操作序列构建为集合通信层流水线；将各计算单元的专用计算引擎归约操作构建为归约流水线；以及通过硬件启动函数将拷贝流水线、集合通信层流水线和归约流水线部署到对应的硬件引擎上并行执行。

11、在一些实施例中，在人工智能加速器集群中，每个计算单元包括独立的dma引擎和专用计算引擎，主计算单元通过片内高速互联与同卡中的从计算单元通信，主计算单元通过pcie交换机与其他卡的主计算单元通信，不同节点的主计算单元之间通过远程直接存储器存取rdma网络通信。

12、根据本公开的一方面，提供了一种用于人工智能加速器集群的集合通信装置，人工智能加速器集群中的每张加速卡包括多个计算单元，该装置包括：第一模块，被配置为从每张加速卡中的多个计算单元中选取一个计算单元作为主计算单元，除主计算单元之外的至少一个计算单元作为从计算单元；第二模块，被配置为基于所有加速卡的主计算单元构建单向逻辑环作为环维度，基于同一加速卡上的所有计算单元构建包括父子关系的组作为组维度；第三模块，被配置为根据每个主计算单元在环维度上的位置以及该主计算单元是否具有跨卡或跨机网络通信能力，对环维度上的各个主计算单元进行分类；第四模块，被配置为生成集合通信操作的通信执行序列，通信执行序列包括多个按时间顺序排列的硬件可执行的操作指令，操作指令至少包括dma传输指令、同步事件轮询指令、同步事件提交指令和专用计算引擎归约指令；以及第五模块，被配置为基于通信执行序列中的操作指令构建为多个并行流水线，并部署到人工智能加速器集群的硬件上执行，其中，对于每个通信轮次，通信执行序列指示每个主计算单元在环维度上与该主计算单元的前驱主计算单元和后继主计算单元执行跨卡数据收发，同时在同一通信轮次内，通信执行序列指示同一加速卡内的计算单元在组维度上沿父子关系执行片内数据归约和复制，使得环维度的跨卡通信和组维度的片内通信在时间上至少部分重叠地并行执行。

13、根据本公开的另一方面，提供了一种电子设备，包括处理器；以及存储程序的存储器，程序包括指令，指令在由处理器执行时使处理器执行本公开如上所提供的方法的步骤的方法。

14、根据本公开的另一方面，提供了一种计算机可读存储介质，其上存储有计算机程序，计算机程序被处理器执行时，使处理器执行本公开如上所提供的方法。

15、根据本公开的另一方面，提供了一种计算机程序产品，包括计算机程序，该计算机程序被处理器执行时，使处理器执行本公开如上所提供的方法。

16、根据本公开的一个或多个实施例，通过将跨卡通信作为环维度以及将片内通信作为组维度进行分层处理，充分利用了ai加速器的多计算单元架构优势。在环维度上，只有主计算单元参与跨卡/跨机的环上数据传输，减少了跨卡链路上的数据流量；在组维度上，同卡内的多个计算单元通过片内高带宽互联进行数据归约和同步，充分利用了片内互联的高带宽（通常远高于跨卡pcie带宽）。这种层次化设计避免了所有计算单元都经由跨卡链路通信的低效模式，实现了高效的多计算单元集合通信。

17、根据在下文中所描述的实施例，本公开的这些和其它方面将是清楚明白的，并且将参考在下文中所描述的实施例而被阐明。