\

IPU程序的优化之一:初识内存布局

作者:

分享  

Share on weixin
Share on weibo
Share on linkedin

订阅

内存布局(memory layout)为什么重要

现代处理器,计算速度远比内存的读写速度快。因此,不仅仅是IPU,其他各种处理器,都对内存布局有一定的要求,只有内存的布局满足一定的条件,才能充分发挥出应有的算力。不同的是,内存的布局在不同的处理器上有不同的表现形式。以GPU为例,在神经网络中,我们通常会将tensor的格式定义成NCHW,考虑到GPU算力的发挥,在数据加载指令数一定的情况下,我们需要追求更多的计算指令数,即前者和后者的比值越低,算力的发挥越充分。同时,考虑到连续访存,以及cache映射的效率,会尽可能的将tensor的格式定义成NC4HW4(或相似结构,比如NHC4W4)。

1 NCHW到NC4HW4的转换

这部分内容涉及到代码示例,以及一些性能分析,受篇幅限制,本文不再赘述,感兴趣的话可以参考这个链接了解更多。

这种NC4HW4就可以理解成一种内存布局,它对于算力的发挥至关重要,可以说,它是后续一切优化工作的起点。

本文后续章节会反复提到PopVision。PopVision是Graphcore为开发者提供的一套性能诊断工具,对于分析一段IPU程序是否高效有着重要的作用,如欲了解更多详情,请访问这个链接

什么是IPU的内存布局

内存布局对于充分发挥IPU的算力性能有特别重要的意义。首先,我们需要清楚什么是IPU上的内存布局。

以MK2为例,一个IPU有1472个tile,每个tile上有624K的内存(片上),一共约900M的内存,如图2所示。

图2 MK2内部架构

现在,我们假设有一个float32的tensor,其shape是[64, 128, 256]。很显然,一个tile无法放下这么多的数据,必须将这个tensor的数据以合理方式分配到各个tile上,其分配的结果就是我们要说的内存布局。

以传统的CPU/GPU为例,其内存是中央控制式的,即一颗芯片可以访问/共享全部的内存空间。相反,在IPU内部,每个tile都是独立的,因此IPU对内存的访问是标准的分布式控制,即每个tile各有其独立的内存空间。在IPU程序中,若要跨tile对内存进行访问,就要涉及到tile间的数据交换,这样就会产生一定的消耗。如何尽可能降低这部分的消耗,也是我们需要重点关注的问题。

我们以一个简单的问题y=1-x为例,通过PopVision,观察一下IPU程序都有哪些消耗:

图3 IPU程序的消耗

从图3可以看出,为了执行一个简单的程序,IPU需要执行这几个部分:

  1. DoExchange,tile间的数据交换;
  2. PreArrange, 针对输入数据,tile内部对内存的布局进行重排;
  3. DeExecute, 实际要执行的程序;
  4. PostArrage,如有需要,针对输出数据,tile内部对内存的布局进行重排。

除了实际要执行的程序,其他的消耗都是由内存布局引起的,从上图可知,内存布局对IPU程序的重要性。

为了降低由内存布局引起的各种消耗,通常我们需要从两个方面考虑问题:

  1. 以原地访存的方式,合理地在tile上分配内存和vertex。这样可以避免tile间的数据交换;
  2. 避免不连续的内存的拷贝,比如transpose/subsample/reverse等等,这也是为了避免tile间的数据交换以及一些tile内部的数据重排。

TileMapping

对于一个给定的计算任务,在IPU编程模型中,我们需要将输入数据、计算程序,和输出数据合理地映射(或称分配)到各个具体的tile上,这里的“映射”就是TileMapping。

图2描绘了IPU程序需要将哪些对象映射到具体的tile上,其中“Tensors”(包括t1和t2)表示输入数据,“Vertices”(蓝色部分的圆圈)表示要执行的计算程序,最后的t3和t4表示输出数据。

图4 IPU编程模型示意

结合图4,关于内存布局我们需要清楚以下几点:

  1. 输入数据(input tensor),实际程序中,输入数据的布局是由上一个op决定的,或者是在程序开始处决定。作为op的开发者,我们无法确定输入数据的布局是什么样的。如果按照开发者的期望,对输入数据进行TileMapping,就会造成一个tensor的多重mapping,导致同一段数据被多个tile引用,会对程序造成极大的隐患,这点在后文将有详细的解释;
  2. 计算程序(vertex),vertex如何mapping取决于op的算法特性和开发者的设计,因此,由开发者决定vertex的TileMapping;
  3. 输出数据(output tensor),输出数据的映射由op的开发者决定,其映射的结果就是输出数据(下一个op的输入数据)的布局。

基于以上三点,为了合理规划vertex和输出数据的布局,如何确定输入数据的布局就非常重要,我们需要根据输入数据的布局来确定如何将vertex和输出数据映射到相应的tile上。

现在解释一下上述第一条,同一段数据被多个tile引用,会有什么不利影响呢?我们可以看一个例子。假设有一个float32的tensor,命名为A,其shape是[64, 32768],要经过两个op,假设这两个op都仅仅做数据copy。初始化的时候,A按行进行TileMapping,这样就会映射到64个tile上(tile0-tile63)。我们再看一下两个op对其输入输出的映射步骤,正常情况下,第一个op是这样的:

Tensor B = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “B”); //创建tensorB

for(int i= 0 ; I < 64; i ++)

  graph.setTileMapping(B[i], i);

prog.add(poplar::program::Copy(A, B)); //指定A往B进行copy

然后,第二个op是这样的:

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

for(int i= 0 ; i < 64; i ++)

  graph.setTileMapping(C[i], i);

prog.add(poplar::program::Copy(B, C));//指定B往C进行copy

为了说明什么是多重映射,我们假设在第二个op中对其输入(也就是B)进行了多重映射,假设又把B映射到了tile64-tile127,那么其代码像这样:

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

for(int i= 0 ; i < 64; i ++)

{

  graph.setTileMapping(B[i], i + 64);

  graph.setTileMapping(C[i], i);

}

prog.add(poplar::program::Copy(B, C)). //指定B往C进行copy

最后,我们需要通过PopVision观察一下二者的区别,如图5、图6所示:

图5 没有多重mapping的Tile Memory
图6 tile64-tile127进行的第二重mapping

从上图可以看到,B被多重(准确地说是2重)映射,在tile0-tile63和tile64-tile128都有一份副本,造成了内存的浪费。为了说明其危害,我们还可以做这样的假设:假设这个多重映射就发生在tile0-tile63,我们在B的高维度做transpose(仅仅是view层面)这样的操作。

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

Tensor BT = B.reshape({ 8, 8, 32768 });

BT = BT.dimShuffle({ 1, 0, 2 }); //在view层面,高维度进行transpose

BT = BT.reshape({ 64, 32768 })

for(int i= 0 ; i < 64; i ++)

{

  graph.setTileMapping(BT[i], i);

  graph.setTileMapping(C[i], i);

}

prog.add(poplar::program::Copy(B, C)). //指定B往C进行copy
图7 形变后的多重mapping

综合图5、图6、图7可以看出,开发者在实践中要尽量避免这种多重映射的发生。

原地访存

原地访存,很容易联想到“inplace”。本文所述的“原地访存”跟通常所说的“inplace”是有区别的。一般而言,inplace是指,一段内存既是输入也是输出。本文所述的“原地访存”是指在原有输入数据的布局基础上对数据进行读取。

举个例子,假设有一个一维的tensorA,其shape是[32],其中0-15在tile0上,16-31在tile1上,vertex在tile2上。显然,vertex跟数据不在同样的tile上,其对数据的读取必然涉及到tile间的数据交换,这种方式就不是原地访存。

相反,如果在tile0上配置一个vertex读取tensorA的0-15,在tile1上配置一个vertex读取tensorA的16-31,那么这种数据的读取方式就是原地访存。因为vertex和tensorA的各个切片是在相同的tile上的,这样就不会产生tile之间的数据交换,对于IPU程序的性能发挥有很大的帮助。

假设tensor是以NCHW来定义的。在许多程序中,数据的处理会不自觉地按照最后一个维度进行的,即一行一行地处理。这样我们就会按照逐行的方式确定vertex和输出数据的布局。1.3中已经说明了输入数据的布局是由上一个op决定的,开发者无法确定,输入数据有可能是这样的:

  1. 布局不是按照逐行的方式;
  2. 不是一个tensor实体,可能是由不同的tensor拼接起来的;
  3. 内存不是连续的。

如果op开发者仍然按照自己期望的逐行的方式确定自己的vertex和输出数据的布局,那么就会产生一定的消耗。举个例子,假设有A和B两个tensor,每个tensor的shape都是[32, 64, 128, 128]。每个tensor都按照线性的方式将数据均匀地映射到每个tile上,然后两者在最后一个维度进行拼接(concat),最终形成tensorC。那么,这个C的shape就是[32, 64, 128, 256]。我们假设按照逐行的方式对数据进行处理(计算任务就是简单的y=1-x),即vertex和输出数据都是按照逐行的方式进行tile的映射的。最后,我们通过PopVision,观察其执行的情况,如图8。

图8 没有原地访存的情况
Cycles(DoExchange + PreArrange + Execute) = 117448

很显然,tile间的数据交换和tile内部的PreArrange占据了整个程序的大部分执行时间,实际程序的执行时间不到整个程序执行时间的一半,这造成了算力的极大浪费。

实际上,如果我们能够充分利用输入数据的布局来规划vertex和输出数据的布局,这些浪费是可以避免的。事实上,像y=1-x这样的计算任务,我们确实可以达成原地访存的目的。IPU编程的基本接口提供了一些关于分析tensor布局的接口,具体详情可以参考这个链接,这里简单介绍一下:

  1. graph.getTileMapping:获取待定tensor在每个tile上的view层面的数据分布;
  2. graph.getSortedContiguousRegions:根据上述view层面的数据分布,获取每个tile上物理上的分布,并且按照连续区块分组;
  3. graph.reorderToSimplify:重排以简化待定tensor的view层面的结构,尽可能将物理上的连续区域合并,这个接口一般对于顺序无关(比如elementwise)的操作有用,本例即可适用。

以上述引用的例子为例,为了达成原地访存的目的,我们可以利用这些接口做一些准备工作:

  1. 遍历每个tile,通过上述接口获取每个tile上关于目标tensor的每个切片的分布情况;
  2. 根据每个输入切片在每个tile上的分配情况,分配相应的vertex;
  3. 根据每个输入切片在每个tile上的分配情况,分配相应的输出切片;
  4. 关键在于第一点,如何获取到每个切片的分配信息。这个程序相对复杂,本文不再展开,具体可以参考本文对应的示例代码。

最后,我们还是通过PopVision中profile的情况,观察一下优化后的结果,如图9所示。

图9 引入原地访存的情况
Cycles(Execute) = 45890

很明显,DoExchange和PreArrange都消失了,只有纯计算部分。整个程序的性能提升了1倍以上。

利用现有tensor的布局,进而按照原地访存方式进行数据的处理,优势是很明显的,但同时也有相应的劣势:原地访存的方式会导致tile间工作负载的不均衡。举个在实际中可能会用到的例子:

现有tensorA和B,A的shape是[2, 262144, 64],B的shape是[2, 64, 64],

令C=matmul(A,B),O=C[0] + C[1]

我们观察最后一步O=C[0] + C[1],这一步的工作负载,如图10所示。

图10 tile负载不均衡

C=A+B这种ElementWise的操作属于典型的原地访存的解决方式。这种情况下,IPU程序会以其中一个tensor的布局为准进行原地访存的处理。显然,在这种情况下,会产生工作负载不均衡的问题。如图6所示,有一半的tile被闲置了。

因此,原地访存有其优势,也有其劣势,需要根据实际情况决定是否适用。

最后附上本文用到的例子的源代码的链接:

https://phabricator.sourcevertex.net/diffusion/PUBLICEXAMPLES/browse/ipu_performance_test/code_examples/poplar/ipu_performance_test/

More Posts

ChatGPT开源平替:OpenAssistant OASST1微调版Pythia-12B

Flan-T5:用更小且更高效的LLM实现出色效果

详细攻略:在IPU上以float16精度运行FLAN-T5-XL推理

较小模型,超高性能:DeBERTa和自然语言理解的未来

PackedBert:如何用打包的方式加速Transformer的自然语言处理任务

Pienso为企业提供由云上IPU支持的高效大型语言模型访问

获取最新的GRAPHCORE资讯

在下方注册以获取最新的资讯和更新:




    获取最新的GRAPHCORE资讯

    在下方注册以获取最新的资讯和更新: