昇思MindSpore学习入门-静态图高级编程技巧

如何优化编译性能

使用lazy_inline装饰器

神经网络模型的编译过程往往采用默认inline的方式,把层级的代码表达最终展开成一张扁平的计算图,一方面寻求最大的编译优化机会,另一方面也可以简化自动微分以及执行的逻辑。inline后形成的计算图包含了所有的计算节点,可以在更大的范围内进行优化,比如常量折叠、节点融合、并行分析等,也可以更好地实现内存分配,减少内存申请和性能开销。虽然inline优化对于运行期性能提升帮助非常大,但过度inline也带来了编译期的负担。例如随着计算图节点数量膨胀,执行pass的耗时也在急剧增长。

为了减轻inline对编译性能带来的损耗,对于重复调用相同计算单元的场景(典型的场景是在for循环中调用同一个Cell类的不同实例),我们提供了Lazy Inline机制来减少编译时间。

大模型pipeline并行场景

在大模型场景中,编译耗时问题尤为突出,一是大模型的模型结构层次深,节点数多;二是大模型在训练时,由于启用pipeline并行,导致模型规模和节点数进一步加大,如果原来图的规模是O,那开启pipeline并行,单节点图的规模变为(O/X)*Y,其中X为pipeline的stage数量,Y为micro batch的数量。以盘古13B网络为例,计算图中计算节点数量达到13.5万个,单次编译时长可接近3小时。

我们观察到类似盘古的大模型网络结构,是由多层layer组成的,在开启pipeline并行时,各个micro batch的layer层结构是完全一样的。当开启pipeline并行时,PipelineCell使用for循环的方式来多次调用相同结构的layer,代码如下所示:

from mindspore import nn

class PipelineCell(nn.Cell):

def init(self, network, micro_size):

...

self.network = network

self.micro_size = micro_size

...

def construct(self, ...):

...

for i in range(self.micro_size):

output = self.network(...)

...

如果我们把循环体看作被频繁调用的子图,通过把它标记为Lazy Inline,告知编译器推迟inline处理,那么就可以在编译的大部分阶段大幅度减少计算图节点数量,从而获得性能收益。例如上面的代码,可以保留network实例的子图结构,不inline或者不提前inline。对此,我们提供了@lazy_inline装饰器来实现延迟inline。

以Pangu_alpha网络为例,PipelineCell函数体中处理的network为PanGUAlphaWithLoss类的实例,为实现延迟inline,我们需要对PanGUAlphaWithLoss类的__init__函数加上@lazy_inline装饰器,以标记PanGUAlphaWithLoss类的子图结构需要被保留下来,不做inline或者延迟inline。如下所示:

from mindspore import nn

from mindspore import lazy_inline

class PanGUAlphaWithLoss(nn.Cell):

@lazy_inline

def init(self, ...):

...

def construct(self, ...):

更加泛化的一般场景

@lazy_inline是Cell::__init__的装饰器,它会以__init__的所有参数生成Cell的cell_init_args属性值,cell_init_args值相同表明Cell类名和初始化参数值是一样的。而对于相同Cell类的实例,它们的weights还可能是不一样的,因此对于用construct(self, x)定义的网络结构,在实际编译时我们可以转换为construct(x, self.cell_init_args, self.trainable_parameters())。对于同一个Cell类的不同实例,如果cell_init_args是相同的,那么这两个实例可以复用同一个网络结构,如下所示:

def construct(self, x)

reuse_construct(x, self.trainable_parameters())

引入可复用计算图后,具有相同cell_init_args的Cell实例只需编译解析一次。所以对于更加泛化的调用同一个Cell类的不同实例的场景,只要cell_init_args是相同的,我们都可以加上@lazy_inline装饰器来加速编译。例如GPT网络:

from mindspore import nn

from mindspore import lazy_inline

class Block(nn.Cell):

@lazy_inline

def init(self, config):

...

def construct(self, x, attention_mask, layer_past):

...

class GPT_Model(nn.Cell):

def init(self, config):

...

for i in range(config.num_layers):

self.blocks.append(Block(config))

...

self.num_layers = config.num_layers

def construct(self, input_ids, input_mask, layer_past):

...

present_layer = ()

for i in range(self.num_layers):

hidden_states, present = self.blocks[i](...)

present_layer = present_layer + (present,)

...

GPT的网络结构由多层Block类的不同实例构成,这些Block的初始化参数都是同一个config,所以加上@lazy_inline装饰器后,这些Block实例都可以复用同一个网络结构,而且在大部分的编译阶段都不进行inline,从而可以大幅度减少编译时间。

使用步骤

如上面的例子,在网络脚本中,往需要延迟inline和复用子图结构的Cell类的__init__函数加上@lazy_inline装饰器。

使用限制

  1. Cell 是以Cell的类名和__init__参数值生成Cell实例标识的,这是基于__init__的参数确定Cell 的所有属性,以及construct构图开始时的Cell属性和__init__执行完的属性一致为假设前提,因此Cell与构图有关的属性,在__init__执行完后不能进行更改。例如:

from mindspore import nn

from mindspore import lazy_inline

class Block(nn.Cell):

@lazy_inline

def init(self, ...):

self.x = 0

...

def construct(self, ...):

if self.x == 0:

...

else:

...

...

class Model(nn.Cell):

def init(self, ...):

...

self.num_layers = 10

for i in range(self.num_layers):

self.blocks.append(Block(...)) # 此处Block进行初始化

...

self.blocks[0].x = 1 # 此处在Block初始化后修改Block的属性,会导致该Block无法复用同一份子图

def construct(self, ...):

...

for i in range(self.num_layers):

res = self.blocks[i](...)

...

如上代码所示,网络Model中的某个Block实例,它的属性x在该实例初始化后被修改了,那么这个Block实例就无法准确复用同一个子图结构了。

  1. 一个Cell类的网络结构包含多个Cell_X类的实例,同时每个Cell_X类的网络结构又包含多个Cell_Y的实例的场景,如果往Cell_X和Cell_Y类的__init__函数上都加上@lazy_inline,那么只有最外层的Cell_X实例的网络结构被编译成可复用的计算图且被延迟inline,内层的Cell_Y实例的计算图还是会被inline。例如:

from mindspore import nn

from mindspore import lazy_inline

class InnerBlock(nn.Cell):

@lazy_inline # InnerBlock不会被延迟inline

def init(self, ...):

...

def construct(self, ...):

...

class OuterBlock(nn.Cell):

@lazy_inline # OuterBlock将会被延迟inline

def init(self, ...):

...

self.num_layers = 10

for i in range(self.num_layers):

self.blocks.append(InnerBlock(...))

def construct(self, ...):

...

for i in range(self.num_layers):

res = self.blocks[i](...)

...

class Model(nn.Cell):

def init(self, ...):

...

self.num_layers = 10

for i in range(self.num_layers):

self.blocks.append(OuterBlock(...))

def construct(self, ...):

...

for i in range(self.num_layers):

res = self.blocks[i](...)

...

使用HyperMap

使用场景:使用HyperMap替换for循环来优化编译性能。

HyperMap是一个特殊的类,类对象构造时需要传入映射函数f,调用对象时需要传入f的n个参数序列,更多使用方法见:HyperMap。映射函数f必须是MultitypeFuncGraph类型, 可参考MultitypeFuncGraph。在使用for循环批量处理列表元素时,可以通过HyperMap等价语义替换来优化网络编译性能。

使用编译缓存

使用场景:在进行训练或者推理时,如果编译依赖的文件未作任何变更,通过使用编译缓存来缩短编译时间。

编译缓存的本质是存储了网络模型的编译中间过程文件,当网络模型不变时,生产的编译中间过程文件也是一样的,因此可以复用上一次编程产生的中间过程文件。

通过设置context中的enable_compile_cache或环境变量MS_COMPILER_CACHE_ENABLE,可以指定是否保存和加载编译缓存,前者优先级更高。

通过设置context中的compile_cache_path或环境变量MS_COMPILER_CACHE_PATH,可以指定MindSpore编译缓存目录,用于存储图和算子编译过程生成的缓存文件,前者优先级更高。

相关推荐
Tester_孙大壮3 分钟前
走进人工智能体
人工智能
CES_Asia3 分钟前
上海政策引领医学前沿创新,CES Asia 2025展望脑机新篇
人工智能·科技·智能手机·智能音箱·智能电视
Senar3 分钟前
机器学习和前端
前端·人工智能·机器学习
四口鲸鱼爱吃盐11 分钟前
Pytorch | 利用SMI-FGRM针对CIFAR10上的ResNet分类器进行对抗攻击
人工智能·pytorch·python·深度学习·机器学习·计算机视觉
柒月的猫14 分钟前
求和(2022蓝桥杯A组试题C)
c语言·算法·蓝桥杯
纪伊路上盛名在16 分钟前
Max AI prompt1
笔记·学习·学习方法
云天徽上18 分钟前
【机器学习案列】车牌自动识别系统:基于YOLO11的高效实现
人工智能·机器学习
simple_ssn22 分钟前
【蓝桥杯】压缩字符串
java·算法
DashVector27 分钟前
如何通过HTTP API插入Doc
数据库·人工智能·http·阿里云·向量检索
我叫czc32 分钟前
【Python高级366】静态Web服务器开发
服务器·前端·python