基于超级计算机的多头注意力机制的前向计算方法及系统转让专利

申请号 : CN202310232055.4

文献号 : CN115952393B

文献日 :

基本信息:

PDF:

法律信息:

相似专利:

发明人 : 刘卫国周正浩段晓辉

申请人 : 山东大学

摘要 :

本发明涉及深度学习技术领域,提供了基于超级计算机的多头注意力机制的前向计算方法及系统,包括:从核在主核中读取权重和数据集中的数据,利用向量化并配合循环展开的方式将数据和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并变成多头表示;从核利用向量化并配合循环展开的方式,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;从核利用向量化并配合循环展开的方式,得到第三结果,并在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置。极大的加速了前向计算过程,达到了超线性加速比,接近了理论加速比。

权利要求 :

1.基于超级计算机的多头注意力机制的前向计算方法,其特征在于,包括:

从核在主核中读取权重和数据集中的数据,并利用向量化并配合循环展开的方式将数据和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并传输回主核,在传输回主核的过程中,将查询矩阵、键矩阵和值矩阵均变成多头表示;

从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵的低两维矩阵乘以多头表示的键矩阵低两维的转置,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;

从核利用向量化并配合循环展开的方式,将第二结果乘以多头表示的值矩阵,得到第三结果,并传输回主核,在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置;

其中,数据集包含三个:第一数据集、第二数据集和第三数据集,第三数据集的数据量最大,第二数据集的数据量居中,第一数据集的数据量最小;对于数据量最小的数据集中的每个数据,每个从核在主核中读取一个批处理;对于数据量最大的数据集中的每个数据,在每个批处理下,采用数据分块的方式,得到若干个矩阵块,每个从核读取一个矩阵块;对于数据量居中的数据集中的每个数据,在每个批处理下,按照序列长度进行等分,得到若干份,每个从核在主核中读取一份;从核在主核中读取数据和权重的过程中,进行主核与从核通信;从核在计算得到查询矩阵、键矩阵和值矩阵的过程中进行从核与从核之间的通信;主核和从核之间采用DMA通信,从核与从核之间采用寄存器通信;在进行寄存器通讯的过程中,第二数据集采用寄存器的行广播通信,第三数据集采用寄存器的行列广播通信;

其中,在将查询矩阵、键矩阵和值矩阵传输回主核的过程中,对于第一数据集和第三数据集,按列写回,并转置;对于第二数据集,按行写回,并转置;

其中,向量化并配合循环展开的方式包括:对矩阵乘最内层for循环进行向量化装载,装载的过程保证装载数据的连续性,全部装载都是连续性数据;并配合循环展开,对向量化过程进行进一步优化。

2.如权利要求1所述的基于超级计算机的多头注意力机制的前向计算方法,其特征在于,从核根据局部数据存储剩余空间读取多头表示的键矩阵或多头表示的值矩阵。

3.如权利要求1所述的基于超级计算机的多头注意力机制的前向计算方法,其特征在于,从核将数据传输回主核时,采用直接内存访问通信。

4.如权利要求1所述的基于超级计算机的多头注意力机制的前向计算方法,其特征在于,对于数据量居中的数据集和数据量最大的数据集,在进行乘法计算时,采用寄存器通信。

5.基于超级计算机的多头注意力机制的前向计算系统,其特征在于,包括:

第一计算模块,其被配置为:从核在主核中读取权重和数据集中的数据,并利用向量化并配合循环展开的方式将数据和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并变成多头表示;

第二计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵的低两维矩阵乘以多头表示的键矩阵低两维的转置,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;

第三计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将第二结果乘以多头表示的值矩阵,得到第三结果,并传输回主核,在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置;

其中,数据集包含三个:第一数据集、第二数据集和第三数据集,第三数据集的数据量最大,第二数据集的数据量居中,第一数据集的数据量最小;对于数据量最小的数据集中的每个数据,每个从核在主核中读取一个批处理;对于数据量最大的数据集中的每个数据,在每个批处理下,采用数据分块的方式,得到若干个矩阵块,每个从核读取一个矩阵块;对于数据量居中的数据集中的每个数据,在每个批处理下,按照序列长度进行等分,得到若干份,每个从核在主核中读取一份;从核在主核中读取数据和权重的过程中,进行主核与从核通信;从核在计算得到查询矩阵、键矩阵和值矩阵的过程中进行从核与从核之间的通信;主核和从核之间采用DMA通信,从核与从核之间采用寄存器通信;在进行寄存器通讯的过程中,第二数据集采用寄存器的行广播通信,第三数据集采用寄存器的行列广播通信;

其中,在将查询矩阵、键矩阵和值矩阵传输回主核的过程中,对于第一数据集和第三数据集,按列写回,并转置;对于第二数据集,按行写回,并转置;

其中,向量化并配合循环展开的方式包括:对矩阵乘最内层for循环进行向量化装载,装载的过程保证装载数据的连续性,全部装载都是连续性数据;并配合循环展开,对向量化过程进行进一步优化。

6.一种计算机可读存储介质,其上存储有计算机程序,其特征在于,该程序被处理器执行时实现如权利要求1‑4中任一项所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。

7.一种计算机设备,包括存储器、处理器及存储在存储器上并在处理器上运行的计算机程序,其特征在于,所述处理器执行所述程序时实现如权利要求1‑4中任一项所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。

说明书 :

基于超级计算机的多头注意力机制的前向计算方法及系统

技术领域

[0001] 本发明属于深度学习技术领域,尤其涉及基于超级计算机的多头注意力机制的前向计算方法及系统。

背景技术

[0002] 本部分的陈述仅仅是提供了与本发明相关的背景技术信息,不必然构成在先技术。
[0003] 自然语言处理,向来是深度学习领域极其重要的研究方向,近些年定制化的AI加速芯片出现和发展极大地提升了计算深度学习框架的能力,自然语言相关的学习框架规模也在不断地扩大,出现了基于变换器的双向编码器表示技术(Bidirectional Encoder Representation from Transformers,BERT)、生成式预训练变换模型2(Generative Pre‑trained Transformer 2,GPT‑2)、生成型预训练变换模型3(Generative Pre‑trained Transformer 3,GPT‑3)等超大规模框架,在机器翻译、语言生成等方面取得了不俗的表现。这些模型大多基于Transformer框架,该框架核心是多头注意力(Multi‑Head Attention)机制,Multi‑Head Attention计算量极大,往往需要庞大的计算资源和极长的计算时间,对其进行深度地优化可以极大缩短训练时间或减少所使用的计算资源。
[0004] 基于深度神经网络的算法,相比于各领域传统算法有更大运算量和参数存储量,因此有更长的计算时延和需要更大的内存空间。针对深度神经网络的参数量和计算量较大的问题,在超级计算机(如神威太湖之光)上,通过异构体组内主从并行可以很好的解决。
[0005] 然而,神威平台上单个核的运行速度并不快,还要低于本地计算机运行速度,但多个核,核组协同计算起来速度会大大提升。如何使多个核进行协同计算,得到优异的效果,是优化的关键;与此同时,在神威太湖之光上对Multihead‑Attention的前向计算过程进行优化,需要在制定并行策略时,进行合理的数据划分,有效提高数据重用率,充分利用从核的局部数据存储(Local Device Memory,LDM);以及进行访存优化,保持写回结果连续性,提高直接内存访问(Direct Memory Access,DMA)的带宽;并且在计算上需要尝试计算的通用优化,减少计算量;并且在优化后还需要使误差不能过大,保证结果的准确性,才能说明优化算法的有效性。

发明内容

[0006] 为了解决上述背景技术中存在的技术问题,本发明提供基于超级计算机的多头注意力机制的前向计算方法及系统,极大的加速了前向计算过程,达到了超线性加速比,接近了理论加速比。
[0007] 为了实现上述目的,本发明采用如下技术方案:
[0008] 本发明的第一个方面提供基于超级计算机的多头注意力机制的前向计算方法,其包括:
[0009] 从核在主核中读取权重和数据集中的数据,并利用向量化并配合循环展开的方式将数据和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并传输回主核,在传输回主核的过程中,将查询矩阵、键矩阵和值矩阵均变成多头表示;
[0010] 从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵的低两维矩阵乘以多头表示的键矩阵低两维的转置,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;
[0011] 从核利用向量化并配合循环展开的方式,将第二结果乘以多头表示的值矩阵,得到第三结果,并传输回主核,在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置。
[0012] 进一步地,对于数据量最小的数据集中的每个数据,每个从核在主核中读取一个批处理。
[0013] 进一步地,对于数据量最大的数据集中的每个数据,在每个批处理下,采用数据分块的方式,得到若干个矩阵块,每个从核读取一个矩阵块。
[0014] 进一步地,对于数据量居中的数据集中的每个数据,在每个批处理下,按照序列长度进行等分,得到若干份,每个从核在主核中读取一份。
[0015] 进一步地,从核根据局部数据存储剩余空间读取多头表示的键矩阵或多头表示的值矩阵。
[0016] 进一步地,从核将数据传输回主核时,采用直接内存访问通信。
[0017] 进一步地,对于数据量居中的数据集和数据量最大的数据集,在进行乘法计算时,采用寄存器通信。
[0018] 本发明的第二个方面提供基于超级计算机的多头注意力机制的前向计算系统,其包括:
[0019] 8、基于超级计算机的多头注意力机制的前向计算系统,其特征在于,包括:
[0020] 第一计算模块,其被配置为:从核在主核中读取权重和数据集中的数据,并利用向量化并配合循环展开的方式将数据和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并变成多头表示;
[0021] 第二计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵的低两维矩阵乘以多头表示的键矩阵低两维的转置,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;
[0022] 第三计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将第二结果乘以多头表示的值矩阵,得到第三结果,并传输回主核,在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置。
[0023] 本发明的第三个方面提供一种计算机可读存储介质,其上存储有计算机程序,该程序被处理器执行时实现如上述所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。
[0024] 本发明的第四个方面提供一种计算机设备,包括存储器、处理器及存储在存储器上并在处理器上运行的计算机程序,所述处理器执行所述程序时实现如上述所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。
[0025] 与现有技术相比,本发明的有益效果是:
[0026] 本发明提供了基于超级计算机的多头注意力机制的前向计算方法,其转变除法为乘法,减少计算的时钟周期,并采取向量化配合循环展开的方式,提高了计算速度,加速了前向计算过程。
[0027] 本发明提供了基于超级计算机的多头注意力机制的前向计算方法,其在将第三结果传输回主核过程中,对主存带跨步写回,设置主存跨步和跨步向量块大小,对数据存储位置进行改变,从而在数据传输过程中达到了数据转置的目的,节省了数据转置,加速了前向计算过程。
[0028] 本发明提供了基于超级计算机的多头注意力机制的前向计算方法,其利用寄存器通信,让核组内数据传输变快,以及良好的主从并行,达到了主从从核工作时间相互掩盖的效果,极大的加速了前向计算过程,达到了超线性加速比,接近了理论加速比。

附图说明

[0029] 构成本发明的一部分的说明书附图用来提供对本发明的进一步理解,本发明的示意性实施例及其说明用于解释本发明,并不构成对本发明的不当限定。
[0030] 图1是本发明实施例一的计算查询矩阵Q、键矩阵K和值矩阵V的流程图;
[0031] 图2是本发明实施例一的Multi‑Head attention流程图;
[0032] 图3是本发明实施例一的寄存器通信的示意图;
[0033] 图4是本发明实施例一的步骤3的流程图;
[0034] 图5是本发明实施例一的步骤4的流程图;
[0035] 图6(a)是本发明实施例一的相乘的两个矩阵中的一个矩阵的示意图;
[0036] 图6(b)是本发明实施例一的相乘的两个矩阵中的另一个矩阵的示意图;
[0037] 图6(c)是本发明实施例一的矩阵块划分示意图。

具体实施方式

[0038] 下面结合附图与实施例对本发明作进一步说明。
[0039] 应该指出,以下详细说明都是例示性的,旨在对本发明提供进一步的说明。除非另有指明,本文使用的所有技术和科学术语具有与本发明所属技术领域的普通技术人员通常理解的相同含义。
[0040] 术语解释:
[0041] 循环展开(Loop unwinding或Loop unrolling):是一种牺牲程序的大小来加快程序执行速度的优化方法。
[0042] Multi‑Head attention:使用输入的X分别和W的转置相乘,计算得到Q、K和V矩阵;然后将Q、K和V矩阵分别变成多头表示Q'、K'和V',即,将D拆分成N×PD的形式,其中PD=D/N,从Q、K和V矩阵中拆出一个头出来,并且把N拿到前面;然后分别将Q'的低两维矩阵乘K'低两维的转置,得到Q'K',即一个批处理(Batch)的矩阵乘;随后采用Norm函数对Q'K'的最低维度执行归一化操作,得到Q"K";之后每个头的Q"K"矩阵乘以V',得到结果Q"K"V';最后,对结果Q"K"V'进行数据转换,将头数N折叠回最后一维中,得到最终结果Y。
[0043] 其中,X为Multi‑Head attention的输入,表示输入编码过后的单词序列,即一个数据,其数据存储布局为[B,S,D],其类型为float;W为Multi‑Head attention的输入,表示Wq、Wk和Wv三个权重连续排列,其数据存储布局为[3,D,D],其类型为float;N为Multi‑Head attention的输入,表示多头的头数,其数据类型为正整数int;B为Multi‑Head attention的输入,表示Batch大小,其数据类型为正整数int;S为Multi‑Head attention的输入,表示序列长度(Sequence),其数据类型为正整数int;D为Multi‑Head attention的输入,表示词向量维度,其数据类型为正整数int;Q、K、V和QK均为Multi‑Head attention的临时存储空间,可用来存储中间结果,其数据类型均为float,Q、K和V的数据存储布局为[B,S,D],QK的数据存储布局为[B,N,S,S];Y为Multi‑Head attention的输出,表示输出的中间特征,其数据存储布局为[B,S,D],其类型为float;QK表示Q'K'或Q"K"。
[0044] Multi‑Head attention的步骤如图2(图2中B=1、S=3、D=4、N=2)所示,具体包括:
[0045] (1)使用输入的X和W的转置分别计算Q、K、V,过程可表示为:
[0046]
[0047]
[0048]
[0049] 其中,Wq、Wk和Wv的最低维度是与X的公共维度,即相乘累加的维度;D'=D/N。
[0050] (2)将Q、K、V转换为多头表示,分别为Q'、K'、V',过程可表示为:
[0051]
[0052]
[0053]
[0054] (3)分别将Q'的低两维乘以(矩阵乘)K'的低两维的转置,得Q'K';然后对Q'K'的最低维度执行归一化(Norm)操作,并除以 ,得到Q"K",过程可表示为:
[0055]
[0056] 其中,Norm的公式如下Norm(Xi)=Xi/∑Xi。
[0057] (4)分别将Q"K"的低两维乘以(矩阵乘)V'的低两维得到结果Q"K"V',并对结果做数据转换,将头数N折叠回最后一维,得到最终的结果Y,过程可表示为:
[0058]
[0059]
[0060] 图2中的初始存储布局左高右低,比如[M,N],表示N维度连续,M为高维,N为低维;W的维度[3,D,D],分别为[Wq,Wk,Wv],Wq、Wk、Wv的维度均为[D,D]。
[0061] 实施例一
[0062] 本实施例提供了基于超级计算机的多头注意力机制的前向计算方法,在保证了结果正确性的前提下,对Multi‑Head Attention过程进行了数据划分、DMA优化、计算优化、寄存器通信和主从并行等优化。
[0063] 在本实施例中,超级计算机为神威太湖之光。
[0064] 本实施例提供了基于超级计算机的多头注意力机制的前向计算方法,包括如下步骤:
[0065] 步骤1、将权重W和数据集输入主核。
[0066] 步骤2、如图1所示,从核在主核中读取权重W和数据集中的数据X,从核利用向量化并配合循环展开的方式将数据X和权重W转置相乘,计算得到查询矩阵Q、键矩阵K和值矩阵V,并传输回主核;然后在传输回主核的过程中将查询矩阵Q、键矩阵K和值矩阵V分别变成多头表示的查询矩阵Q'、多头表示的键矩阵K'和多头表示的值矩阵V',并存储在主核。
[0067] 目前数据划分相关研究工作较少,但其效果却直接决定了整体计算的并行化程度,进而能够影响并行计算的整体性能。一般需要考虑各个从核上的计算负载是否均衡,数据之间有无相关性,并尽可能使其方便提高DMA带宽。
[0068] 本实施例中的数据集包含三个:第一数据集、第二数据集和第三数据集,第三数据集的数据量最大,第二数据集的数据量居中,第一数据集的数据量最小。对于第一数据集中的每个数据X,每个从核在主核中读取一个Batch,每个数据X的存储布局为[B,S,D]。对于第一数据集,一个从核一次只处理S×D,然后主核进行转置,和下一次处理S×D的时间相互掩盖;对于第二数据集中的每个数据X,在每个Batch下,按照序列长度进行等分,得到若干份,每份的大小为(S/8)×D,每个从核在主核中读取一份;对于第三数据集中的每个数据X,在每个Batch下,采用数据分块的方式,得到若干矩阵块,每个从核在主核中读取一个矩阵块。
[0069] 权重划分为N份,等同于多头表示,减少数据转换。对于第一数据集对应的权重W(Wq、Wk和Wv)平均划分为N份后,再将每一份均分到64个从核上,即将第一数据集对应的权重划分为1/(N×64)份;对于第二数据集对应的权重W平均划分为N份后,再将每一份(即每个头的数据(D×pd))分成四份,每个从核读取的三个权重都是D×pd/4,即一列从核对于三个权重来说都是读取两个头的数据;对于第三数据集对应的权重W平均划分为N份后,再将每一份进行矩阵分块,每个从核在主核中读取一个矩阵块。
[0070] 本实施例,根据数据集大小,采取寄存器通信和DMA通信相结合的方式最大化利用DMA带宽和LDM储存,每次都尽可能使LDM空间充分利用。具体的,如图1所示,从核在主核中读取数据时,需要判断数据集大小,根据数据集情况从核进行DMA通信或寄存器通信,具体的:从核在主核中读取数据和权重的过程中,进行主核与从核通信;从核在计算得到查询矩阵Q、键矩阵K和值矩阵V的过程中进行从核与从核之间的通信,对于第一数据集在计算得到查询矩阵Q、键矩阵K和值矩阵V的过程中,不需要进行从核与从核之间的通信。主核和从核之间采用DMA通信,从核与从核之间采用寄存器通信。在进行寄存器通讯的过程中,第二数据集采用寄存器的行广播通信,第三数据集采用寄存器的行列广播通信。
[0071] 在将查询矩阵Q、键矩阵K和值矩阵V传输回主核的过程中,对于第一数据集和第三数据集,按列写回Q、K、V,并转置Q'、K';对于第二数据集,按行写回Q、K、V,并转置V'。这种方式避免了分成多头的步骤且保持了连续写回。
[0072] 可以看出,本实施例对数据计算采取自适应策略,每次计算充分利用从核LDM大小,减少读写次数。具体的,在计算查询矩阵Q、键矩阵K和值矩阵V的过程中,对第一数据集、第二数据集和第三数据集采取不同的优化方式,即,第一数据集、第二数据集和第三数据集采取的通信方式不同:第一数据集直接通过DMA通信,将数据从主核读取到从核,从核LDM空间充足,不需要在从核之间再进行寄存器通信;第二数据集和第三数据集还采取了寄存器通信,进行从核之间数据交互,导致数据顺序有变化,所以计算的时候也不相同,最后写回时数据顺序也不同。
[0073] 在计算查询矩阵Q、键矩阵K和值矩阵V时,对数据进行向量化,在向量化过程中配合循环展开,使向量化过程进行流水线作业,进一步减少计算时间,在此阶段不同数据集采取不同方式,尽最大可能进行优化。
[0074] 步骤3、如图4所示,从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵Q'的低两维矩阵乘以多头表示的键矩阵K'低两维的转置,得到第一结果Q'K';然后,通过转变除法为乘法,并采取向量化配合并循环展开的方式,对Q'K'的最低维度执行归一化操作,得到第二结果Q"K",并传输回主核。
[0075] 从核在计算第一结果Q'K'的过程中,从核在主核中读取查询矩阵Q'和键矩阵K'。即,从核根据数据集的大小以及LDM剩余空间读取查询矩阵Q'和键矩阵K',具体的:从核首先根据数据集(Q')的大小,将Q'的数据完整读取到(若数据集(Q')的大小超出从核存储空间,则先读取一部分Q',再读取一部分Q',最终完整读取到Q');然后,根据从核LDM剩余大小决定读取的K'的大小,进行K'的读取;最后在计算Q'K'时,进行向量化配合循环展开。实现对计算过程进行优化。
[0076] 如图4所示,从核在执行归一化操作的过程中,首先对Q'K'矩阵中的数据进行求和sum的计算,并采取向量化;然后在后续除法计算中,将sum转为倒数1/sum,转变除法为乘法,并采取向量化配合循环展开的方式,提高计算速度。在从核与主核传输数据时,保证了数据的连续性,提高存取速度。
[0077] 步骤4、如图5所示,在得出最终结果的阶段,从核利用向量化并配合循环展开的方式,将第二结果Q"K"乘以多头表示的值矩阵V',得到第三结果Q"K"V',并传输回主核,在将第三结果Q"K"V'传输回(写回)主核时,对主存带跨步写回,改变数据存储位置,完成转置。
[0078] 从核在计算第三结果Q"K"V'的过程中,从核根据数据集的大小以及LDM剩余空间,在主核中读取第二结果Q"K"和多头表示的值矩阵V'。具体的:从核首先根据数据集(Q"K")的大小,将Q"K"的数据完整读取到(若数据集(Q"K")的大小超出从核存储空间,则先读取一部分Q"K",再读取一部分Q"K",最终完整读取到Q"K");然后,从核根据LDM剩余空间决定读取的V'的大小,进行V'的读取;最后在计算Q"K"V'时,进行向量化并配合循环展开,进行加速。
[0079] 在本实施例中,对于第二和第三数据集中的数据在第一阶段,在进行乘法计算时,先采用寄存器通信进行从核间通信,传输数据,然后进行乘法计算。
[0080] 完成第三结果Q"K"V'的计算后,将第三结果Q"K"V'传输回主核,在将第三结果Q"K"V'传输回主核过程中,对主存带跨步写回,设置主存跨步和跨步向量块大小,对第三结果Q"K"V'的数据存储位置进行改变,从而在传输第三结果Q"K"V'的过程中达到了数据转置的目的,节省了一步数据转置。
[0081] 在本实施例中,从核将数据传输回主核时,采用DMA通信,DMA通信是常用的主存与LDM的数据交互方式。本实施例通过主存跨步调整从核计算结果的写回顺序,保证了写回数据的连续性,提高了DMA的带宽;同时,在从核已经读取相同数据时,采用GET的广播模式,提高读取数据的效率,减少DMA操作次数,降低主存和从核组间通信总量,避免从核间竞争访存。
[0082] 在本实施例中,在计算上的优化主要是利用向量化并配合循环展开,充分提高计算性能,使其接近机器峰值;并且将除法操作变化成乘倒数,减少计算的时钟周期。
[0083] 在本实施例中,向量化并配合循环展开的方式包括:对矩阵乘最内层for循环进行向量化装载,装载的过程保证装载数据的连续性,全部装载都是连续性数据,更好的利用了高速缓存(Cache);并配合循环展开,对向量化过程进行进一步优化。
[0084] 循环展开后方便编译器识别相关代码部分做一些自动的优化,同时减少循环次数。循环展开会优化指令流水线,充分利用寄存器。循环展开后,还利用装载寄存器的时间掩盖寄存器计算的时间。于此同时,矢量类型的数据项中,除了floatv4类型是16字节对齐以外,在主存储器内是32字节对;申威26011处理器架构中,无对齐的装载(Load)/保存(Store)发生异常,操作系统接收到异常信号时,这些Load/Store分割为标准类型的Load/Store,性能大幅下降;向量化过程,可以保证所有向量化过程都是对齐的,使性能得到保证。
[0085] 由于从核的局部存储空间有限,申威处理器提供了核组内同行(列)的寄存器级通信机制,为了能够更好的利用片上存储资源、实现从核间的高效信息交互。与DMA方式相比,DMA通信的带宽仅仅只有数十GB/S,而寄存器通信的带宽高达几百GB/S,通信速度更快。
[0086] 在超级计算机中,一个核组包括一个主核和64个从核,64个从核编号为从0到63,分为8行8列排布,第一行为0到7号从核,第二行为8到15号从核,依次类推,第8行为56到63号从核。
[0087] 在本实施例中,按照序列长度进行数据等分的方式具体为:在某一个Batch下,将S×D分成了8份,每一列取(S/8)×D大小的数据,也就是说第一行第一个列从核取到的就是(0到S/8)×D,第一行第二列从核取到的就是(S/8到2×S/8)×D的数据,……,第一行第八列从核取到的就是(S‑S/8到S)×D的数据,8号从核也就是第二行第一列从核上的就是该Batch的下一个Batch的数据依次类推。从核在主核完成数据集的读取后,进行寄存器通信;因为寄存器通信比DMA通信更加迅速,可以节约通信时间,对于第二数据集,在进行乘法计算时通过行广播通信:设置循环,每一行从核依次向其他从核进行通信,首先是第一行从核向其余行从核通信,通信后其余行与接收到的数据进行计算,传输数据的行与自身数据进行计算。即,每次读取8×S×D的数据,一行从核总共读取S×D的数据,从核通过寄存器行广播通信进行传输数据,进行计算,8行从核就能一次处理8×S×D的数据,完成一次读取后,若数据存在剩余,进行下次读取。
[0088] 在本实施例中,数据分块的方式具体为:为能够充分利用从核LDM空间,将矩阵数据X在某一个Batch下,按照核组的大小(即8×8)进行等分,得到若干个矩阵块,每个从核读取一个矩阵块,例如16×16的矩阵,划分后每个从核上就是2×2的矩阵块;于此同时,对于第三数据集,在进行乘法计算时采用寄存器通信行广播和列广播方式,进行数据交互,以完成矩阵乘计算。
[0089] 寄存器通信行广播和列广播方式为:以图3为例,jk代表的是从核在X中读取的数据,ik代表的是从核在W中读取的数据,k=1,2,…,8,在X中读取的数据进行列广播通信,在W中读取的数据进行广播通信,首先将i1和j1的数据传输到从核9,通过向量化配合循环展开进行计算,得出结果ret;接着从核9计算自身的数据,即ret=ret+j2×i2;然后通过寄存器通信,将j3和i3的数据传送到从核9,ret=ret+j3×i3;依次类推,直到ret=ret+j8×i8,就可以计算出一个完整地结果,存在从核9中,然后将数据传回主核。其中,ret相当于用于储存从核9上矩阵计算的结果;i1、j1分别是从核2、和从核8上分到的分割后的矩阵块;i2、j2为从核9上分到的分割后的矩阵块;依次类推,i3和j3是从核10和从核17上分到的分割后的矩阵块;i8和j8是从核15和从核57上分到的分割后的矩阵块。
[0090] 在进行乘法计算时的数据传输过程通过寄存器通信完成,ik、jk代表了两个矩阵被分割后的矩阵块,ik、jk相乘就代表了两个矩阵中的矩阵块相乘,矩阵如图6(a)和图6(b)所示,i代表行,j代表列,ik、jk完整相乘后,就代表了两个矩阵中某些行和某些列的相乘结束,从核有64个,所以要将矩阵数据按照8×8划分,例如16×16的矩阵,划分后每个从核上就是2×2的矩阵块,如图6(c)标出了从核0分到的2×2矩阵块。
[0091] 在本实施例中,主核在从核第一阶段(计算Q、K、V并转化为多头表示的阶段)运行时,主核进行数据预处理,使与从核通信的数据尽可能连续,减少通信开销,因为矩阵分块后,数据不连续,DMA传输连续数据更快,将数据存储位置转化为连续的再传输;将循环按照Batch分割成多段,使从核核组进行计算的同时,主核可以对上一阶段计算的Q、K和V进行转置,来掩盖转置所需的大部分时间;并且,对于第二数据集,最后一次计算完成后,暂不进行V'的转置,先进行Q"K"的计算,在此同时进行最后一次计算出来的V'的转置,完全的对V'的转置过程进行了掩盖。其中,主核进行数据预处理的过程中,也可与从核计算相互掩盖。
[0092] 第一数据集、第二数据集和第三数据集在第一阶段分别采用不同的方法优化,在第二阶段(步骤3)和第三阶段(步骤4),采用相同的方法优化,因为二三阶段消耗时间小,主要时间消耗在第一阶段,采用不同方法优化提升空间不大。
[0093] 例如,初始版本在一个主核上运行,第一数据集具体参数为B的大小为10,S的大小为64,D的大小为384,N的大小为12,该数据集运行时间为1416.285ms;第二数据集具体参数为B的大小为16,S的大小为64,D的大小为384,N的大小为12,该数据集运行时间为2265.628ms;第三数据集具体参数为B的大小为20,S的大小为512,D的大小为768,N的大小为24,该数据集运行时间为121766.029ms;这三个初始数据集运行结果的相对误差都小于
10的负五次方,通过了检验函数。
[0094] 优化版本使用一个异构群,也就是1个主核和64个从核运行,线性加速比为64倍左右,加速比是评价并行算法的重要指标,是度量多核心并行处理相对于单核心处理的加速倍数。
[0095] 对于第一数据集,运行时间为5.397ms,并且通过了检验函数,相对误差低于10的负五次方,加速了238倍,实现了超线性加速比,64个从核加向量化,理论加速比为256倍,除去用于数据传输的开销以及从核仅支持单流水线浮点计算,在理论情况下已经是接近极限的情况。
[0096] 对于第二数据集,运行时间为8.667ms,并且通过了检验函数,相对误差低于10的负五次方,加速了261倍,实现了线性加速比,并且超过了256倍的理论加速比。
[0097] 对于第三数据集,运行时间为581.173ms,且通过了检验函数,相对误差低于10的负五次方,加速了210倍,除去用于数据传输的开销以及从核仅支持单流水线浮点计算,在理论情况下已经是接近极限的情况。
[0098] 以上优化过程均实现了超线性加速比,并且接近理论加速比,其中有一个数据集超过了理论加速比,并且均通过了自定义的检验函数,检验计算值减去正确值的结果的绝对值除以计算值的绝对值与正确值的绝对值的和,是不是大于10的负5次方,如果大于,则没有通过,检验公式为 ,a为计算值,b为正确值,fabs表示取绝对值函数,该方法保证了Multi‑Head Attention的正确性,其检验函数比相对误差更加严格,可以保证相对误差均低于10的负5次方。
[0099] 本实施例在神威太湖之光上针对Multi‑Head Attention进行优化,通过合理的数据划分、访存优化、计算优化、寄存器通信和主从并行的手段,极大的加速了前向计算过程,达到了超线性加速比,接近了理论加速比。具体的,首先通过数据划分,减少了数据之间相关性;然后在DMA传输过程中,尽可能使数据连续,增加DMA带宽,在计算的过程中,利用向量化以及循环展开,计算过程大大提速,并且利用寄存器通信,让核组内数据传输变快,以及良好的主从并行,达到了主从从核工作时间相互掩盖的效果,极大的加速了前向计算过程,达到了超线性加速比,接近了理论加速比。
[0100] 实施例二
[0101] 本实施例提供了基于超级计算机的多头注意力机制的前向计算系统,其具体包括:
[0102] 数据获取模块,其被配置为:将权重和若干个数据集输入主核,主核将权重分为若干份;
[0103] 第一计算模块,其被配置为:从核在主核中读取权重和数据集,并利用向量化并配合循环展开的方式将数据集和权重转置相乘,计算得到查询矩阵、键矩阵和值矩阵,并变成多头表示;
[0104] 第二计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将多头表示的查询矩阵的低两维矩阵乘以多头表示的键矩阵低两维的转置,得到第一结果,并通过转变除法为乘法后,采取向量化配合并循环展开的方式,对第一结果的最低维度执行归一化操作,得到第二结果;
[0105] 第三计算模块,其被配置为:从核利用向量化并配合循环展开的方式,将第二结果乘以多头表示的值矩阵,得到第三结果,并传输回主核,在将第三结果传输回主核时,对主存带跨步写回,改变数据存储位置,完成转置。
[0106] 此处需要说明的是,本实施例中的各个模块与实施例一中的各个步骤一一对应,其具体实施过程相同,此处不再累述。
[0107] 实施例三
[0108] 本实施例提供了一种计算机可读存储介质,其上存储有计算机程序,该程序被处理器执行时实现如上述实施例一所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。
[0109] 实施例四
[0110] 本实施例提供了一种计算机设备,包括存储器、处理器及存储在存储器上并在处理器上运行的计算机程序,所述处理器执行所述程序时实现如上述实施例一所述的基于超级计算机的多头注意力机制的前向计算方法中的步骤。
[0111] 以上所述仅为本发明的优选实施例而已,并不用于限制本发明,对于本领域的技术人员来说,本发明可以有各种更改和变化。凡在本发明的精神和原则之内,所作的任何修改、等同替换、改进等,均应包含在本发明的保护范围之内。