MindSpore技术理解(上)

时间:2024-09-08 12:05:50

MindSpore技术理解(上)

引言

深度学习研究和应用在近几十年得到了爆炸式的发展,掀起了人工智能的第三次浪潮,并且在图像识别、语音识别与合成、无人驾驶、机器视觉等方面取得了巨大的成功。这也对算法的应用以及依赖的框架有了更高级的要求。深度学习框架的不断发展使得在大型数据集上训练神经网络模型时,可以方便地使用大量的计算资源。

深度学习是使用多层结构,从原始数据中自动学习并提取高层次特征的一类机器学习算法。通常,从原始数据中提取高层次、抽象的特征是非常困难的。目前有两种主流的深度学习框架:一种是在执行之前构造一个静态图,定义所有操作和网络结构,典型代表是TensorFlow,这种方法以牺牲易用性为代价,来提高训练期间的性能;另一种是立即执行的动态图计算,典型代表是PyTorch。通过比较可以发现,动态图更灵活、更易调试,但会牺牲性能。因此,现有深度学习框架,难以同时满足易开发、高效执行的要求。

简介

MindSpore作为新一代深度学习框架,是源于全产业的最佳实践,最佳匹配昇腾处理器算力,支持终端、边缘、云全场景灵活部署,开创全新的AI编程范式,降低AI开发门槛。MindSpore是一种全新的深度学习计算框架,旨在实现易开发、高效执行、全场景覆盖三大目标。为了实现易开发的目标,MindSpore采用基于源码转换(Source Code Transformation,SCT)的自动微分(Automatic Differentiation,AD)机制,该机制可以用控制流表示复杂的组合。函数被转换成函数中间表达(Intermediate Representation,IR),中间表达构造出一个能够在不同设备上解析和执行的计算图。在执行前,计算图上应用了多种软硬件协同优化技术,以提升端、边、云等不同场景下的性能和效率。MindSpore支持动态图,更易于检查运行模式。由于采用了基于源码转换的自动微分机制,所以动态图和静态图之间的模式切换非常简单。为了在大型数据集上有效训练大模型,通过高级手动配置策略,MindSpore可以支持数据并行、模型并行和混合并行训练,具有很强的灵活性。此外,MindSpore还有“自动并行”能力,它通过在庞大的策略空间中进行高效搜索来找到一种快速的并行策略。MindSpore框架的具体优势,请查看详细介绍。

深度学习研究和应用在近几十年得到了爆炸式的发展,在图像识别[1]、语音识别与合成[2]、游戏[3]、 语言建模与分析[4]等方面取得了巨大的成功。深度学习框架[5–9]的不断发展使得在大型数据集上训练 神经网络模型时,可以方便地使用大量的计算资源。 目前有两种主流的深度学习框架:一种是在执行之前构造一个静态图,定义所有操作和网络结构, 典型代表是 TensorFlow[5],这种方法以牺牲易用性为代价,来提高训练期间的性能;另一种是立即 执行的动态图计算,典型代表是 PyTorch[6]。通过比较可以发现,动态图更灵活、更易调试,但会牺 牲性能。因此,现有深度学习框架难以同时满足易开发、高效执行的要求。 本文提出了一种新的深度学习框架——MindSpore,该框架旨在实现三个目标:易开发、高效执行 和全场景覆盖。MindSpore 由四个主要组件组成:MindExpression(ME)、GraphEngine(GE)、 MindData(MD)和 MindArmour(MA)。表 1 总结了 MindSpor的技术贡献。

MindSpore技术理解(上)

1. ME 为用户提供了 Python 编程范式,具有以下三个突出特点。

− 自动微分:采用基于源码转换的自动微分机制,在训练或推理阶段,将一段 Python 代码转 换为数据流图。因此,用户可以方便地使用 Python 原生控制逻辑来构建复杂的神经网络模型。

− 自动并行(Automatic Parallelization):由于大规模模型和数据集的不断增加,跨分布式设 备并行化深度神经网络(Deep Neural Network,DNN)训练已经成为一种常见做法。然而, 当前框架(如 TensorFlow[5]、Caffe[10]和 MXNet[7])用于并行化训练的策略仍然简单直接, 而且通常是次优的。MindSpore 并行化训练任务,透明且高效。“透明”是指用户只需更改一行配置,提交一个版本的 Python 代码,就可以在多个设备上运行这一版本的 Python 代码 进行训练。“高效”是指该算法以最小的代价选择并行策略,降低了计算和通信开销。

− 动态图:MindSpore 支持动态图,无须引入额外的自动微分机制(如算子重载微分机制), 从而大大增加了动态图和静态图的兼容性。

2. GE 负责硬件相关的资源管理和优化,将平台特有的信息传递给 ME。GE 接收来自 ME 的数据流图,并将该图中的算子调度到目标设备上执行。GE 将数据流图分解为优化后的子图,并将它们调度到不同的设备上。GE 将每个设备抽象为一个执行引擎(Execution Engine),并提供执行引擎插件机制,用来支持各种不同的设备,而不影响其他组件。

3. MD 负责数据处理,并提供工具来帮助开发者调试和优化模型。通过自动数据加速技术实现了高性能的流水线,以进行数据处理。各种自动增强策略的出现,使用户不必再寻找合适的数据增强策略。训练看板将多种数据集成在一个页面,方便用户查看训练过程。分析器可以打开执行黑匣子,收集执行时间和内存使用的相关数据,从而有针对性地进行性能优化。

4. MA 负责提供工具,以帮助开发者防御对抗性攻击,实现隐私保护的机器学习。在形式方面,MA 提供了以下功能:生成对抗代码、评估模型在特定对抗环境中的性能、开发出更健壮的模=型。MA 还支持丰富的隐私保护能力,如差分隐私[11]、机密人工智能计算[12]、可信协同学习[13、 14]等。

本文的其余内容如下:第 2 章介绍了 MindSpore 的架构;第 3 章通过阐述自动微分、自动并行和动态图支持的设计细节,介绍了 MindSpore 的核心模块——ME;第 4 章介绍了 GE 的工作过程;第 5 章和第 6 章分别介绍了 MD 和 MA 的相关细节;第 7 章介绍了 MindSpore 支持的端云协同架构;为了说明高效执行这一特性,第 8 章阐述了使用 ResNet50 作为基准来评估“自动并行”特性,还介 绍了 MindSpore 的训练和推理性能;最后,第 9 章对我们的工作进行了总结,并介绍了未来的重点 研究方向。

2 MindSpore 概述

2.1 MindSpore 架构如图 1 所示,MindSpore 由 ME 和 GE 组成。

MindSpore技术理解(上)

ME 提供 Python 接口和自动微分功能,可分为前端和后端。前端用于定义用户级应用软件编程接口(Application Programming Interface,API),用于构建和训练神经网络。由于 MindSpore 采用基于源码转换策略的自动微分机制,因此用户可以 Python 方式编程。一种典型的情况是,用户可以用常规的 Python 方式构建复杂的模型和流程控制,如 if、else、while 等。 ME后端是高性能执行和自动微分的核心,自动微分基于源码转换方式。如果用户选择以 Pynative 模式运行,则逐个下发算子运行。如果模型以图模式运行,则后端会使用流水线根据 Python 代码生成计算图,通过解析 Python 代码,生成抽象语法树(Abstract Syntax Tree,AST),然后将其转换 为 A-Normal-Form(ANF)图[15]。ANF 是图形化的,而不是语法化的,因此从算法上操作起来要容 易得多[16]。如果用户需要训练神经网络,流水线自动生成反向计算节点,并添加到 ANF 图中。流水线在构造完整图之后进行许多优化(例如内存复用、算子融合、常数消除等)。如果用户需要在分布式环境中训练模型,流水线将应用自动并行提供的优化策略(参见 3.2 节)。后端的虚拟机通过会话 管理计算图,调用 GE 来运行图,并控制图的生命周期。 GE 负责控制从 ME 接收的数据流图的执行,GE 也是算子和不同硬件资源的管理器。在 Pynative 模式,GE 一次只能从 ME 接收一个算子,并立即将其下发运行。在图形模式下,GE 从 ME 接收一 个计算图,图编译器根据不同的执行引擎,将计算图分解为若干个子图,并对不同的执行引擎采用 不同的优化函数。这一过程充分利用了底层的硬件资源。然后,编译器静态分配必要的资源,并构建将交付到不同硬件资源的内核。任务执行器是内核执行前的最后一步,它通过调用由将运行内核 的相应硬件提供的执行引擎,启动许多任务,以运行接收到的内核。 MD 中的数据处理在训练过程中完成数据流水线,包括数据加载、数据论证、数据转换等。MD 也提 供简单的 API,支持 CV/NLP/GNN 等全场景的数据处理能力。同时,在这个过程中,如何提高数据处理能力,以匹配人工智能芯片的算力,是保证人工智能芯片发挥极致性能的关键。MD 中的 MindInsight 有四个模块:训练看板、溯源、分析器和调试器。分析器可以打开执行黑匣子,以收集执行时间和内存使用的相关数据。调试器是图执行模式下的调试工具,用户可以在训练过程中,查看图的内部结构和节点的输入/输出。MindInsight 对训练过程生成的日志文件进行分析。开发者可以轻松地可视化训练过程,在图形用户 界面(Graphical User Interface,GUI)上与 MindInsight 对比不同的路径。 MA 帮助用户开发更健壮的模型,保护用户在训练和推理数据中的隐私。MA 中的对抗性攻击防御有 三个主要模块:攻击、防御和评估。攻击模块是一个在黑盒和白盒攻击场景下生成对抗样本的工具。 防御模块从攻击模块接收对抗样本,并在训练过程中用这些样本来提高模型的鲁棒性。评估模块提供了多种评估指标,允许开发者更容易可视化其模型的鲁棒性。MA 为了实现隐私保护机器学习,实现了一系列差分隐私感知优化器,这些优化器在训练过程中自动将噪声添加到生成的梯度中。

2.2 编程范式 MindSpore

为用户提供 Python 编程范式。借助基于源码转换的自动微分,用户可以使用原生 Python 控制语法和其他一些高级 API,如元组(Tuple)、列表(List)和 Lambda 表达。为避免用户混淆,MindSpore 引入了尽可能少的接口和概念。在单机平台上训练简单的神经网络时, 用户只需要了解以下五个组件:

张量(Tensor):是一个包含相同数据类型元素的多维矩阵。与其他一些训练框架不同,MindSpore 中没有标量(Scalar)变量(Variable)的概念。要计算一个张量的梯度,该张量的 requires_grad 自动微分属性应该设置为 True。可以用 Numpy 初始化张量,也可以将张量的值转换为一个 Numpy 对象。

数据集(Dataset):是一个单独的异步流水线,该流水线为在训练中将张量无延迟地馈入网络的剩 余部分做准备。

算子(Operator):是构成神经网络的基本计算单元。MindSpore 支持大多数常用的神经网络算子 (如卷积、批标准化、激活等)和数学算子(如加法、乘法等)。此外,MindSpore 支持自定义算子, 允许用户添加新的算子来适配特定的硬件平台,或者合并现有算子来生成复合算子。

单元(Cell):是张量和算子的集合,是所有神经网络单元的基本类。一个单元也可以包含其他单元, 允许它们嵌套在同一个树状结构中。用户通过在单元中定义 construct 函数来表达神经网络的计算逻辑。construct 函数中的计算将在每个调用中执行。

模型(Model):是 MindSpore 中的一个高级 API,它封装了一些低级 API,从而使用户能够更加方便地使用推理和训练功能。如果用户熟悉低级 API,并希望对计算过程建立细粒度的控制,则不一定需要使用模型。 从用户的角度来看,用 MindSpore 编写程序的核心是构建一个对应于神经网络的单元。这个过程首 先从输入张量开始,可以是常量张量或参数张量。然后,用户使用 MindSpore 提供的不同算子构造 一个单元。最后,用户可以使用模型封装这个单元来训练神经网络,也可以直接将输入数据传递给 单元来执行推理任务。

代码 1 展示了 Python 中 MindSpore 程序。

该程序展示了定义以及训练 LeNet[17]神经网络的过程。 代码前 6 行,导入 MindSpore 合适的库。第 7 行到第 25 行定义了与 LeNet 神经网络相对应的 LeNet5 单元。_init_函数实例化 LeNet 使用的所有算子。construct 函数定义了 LeNet 的计算逻辑。第 26 行和第 27 行从 Mnisit 数据集读取数据,并生成一个迭代器 ds 用作训练的输入。第 28 行将 LeNet5 类实例化为 network。用 SoftmaxCrossEntropyWithLogits 函数来计算损失(loss)(第 29 行),并 用 momentum 来优化参数(第 30 行到第 31 行),loss 和 optimizer 用来创建模型(Model)对象。 最后,用 epoch 来控制迭代次数,调用模型的训练方法,并在每个 eval_step 对模型进行评估。

代码 1 LeNet5 的 MindSpore 实现

1 import mindspore.nn as nn

2 from mindspore.ops import operations as P

3 from mindspore.network.optim import Momentum

4 from mindspore.train import Model

5 from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits

6 import mindspore.dataset as de

7 class LeNet5(nn.Cell):

8 def __init__(self):

9 super(LeNet5, self).__init__()

10 self.conv1 = nn.Conv2d(1, 6, 5, pad_mode='valid')

11 self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')

12 self.fc1 = nn.Dense(16 * 5 * 5, 120)

13 self.fc2 = nn.Dense(120, 84)

14 self.fc3 = nn.Dense(84, 10)

15 self.relu = nn.ReLU()

16 self.max_pool2d = nn.MaxPool2d(kernel_size=2

17 self.flatten = P.Flatten()

18 def construct(self, x):

19 x = self.max_pool2d(self.relu(self.conv1(x)))

20 x = self.max_pool2d(self.relu(self.conv2(x)))

21 x = self.flatten(x)

22 x = self.relu(self.fc1(x))

23 x = self.relu(self.fc2(x))

24 x = self.fc3(x)

25 return x

26 ds = de.MnistDataset(dataset_dir="./MNIST_Data")

27 ds = ds.batch(batch_size=64)

28 network = LeNet5()

29 loss = SoftmaxCrossEntropyWithLogits()

30 optimizer = nn.Momentum(network.trainable_params(),

31 learning_rate=0.1, momentum=0.9)

32 model = Model(network, loss, optimizer)

33 model.train(epoch=10, train_dataset=ds)

3 MindExpression

3.1 基于源码转换的自动微分

目前主流的深度学习框架有三种自动微分技术:

l 基于静态计算图的转换:在编译时将网络转换为静态数据流图,然后将链式规则转换为数据流图,实现自动微分。

l 基于动态计算图的转换:以算子重载的方式记录前向执行时网络的操作轨迹,然后将链式规则,应用到动态生成的数据流图中,实现自动微分。

l 基于源码的转换:该技术是从函数式编程框架演化而来,对中间表达(程序在编译过程中的表 达形式),以即时(Iust-In-Time,JIT)编译的形式进行自动微分变换,支持复杂的流程控制场 景、高阶函数和闭包。基于源码转化的自动微分如图 2 所示。 TensorFlow 早期采用静态计算图,而 PyTorch 采用动态计算图。静态图可以利用静态编译技术优化网络性能,但是组建或调试网络非常复杂。使用动态图非常方便,但很难在性能上达到极限优化。

MindSpore 开发了一种新的策略,即基于源码转换的自动微分。

一方面,它支持流程控制的自动微分,因此构建像 PyTorch 这样的模型非常方便。

另一方面,MindSpore 可以对神经网络进行静态编译优化,从而获得良好的性能。

MindSpore 自动微分的实现可以理解为对程序本身进行符号微分,因为 MindSpore IR 是函数式的中间表达,它与基本代数中的复合函数有直观的对应关系,只要已知基础函数的求导公式,就能推导出由任意基础函数组成的复合函数的求导公式。MindSpore IR 中每个原语操作可以对应为基础代数中的基础函数,这些基础函数可以构建更复杂的流程控制。

MindSpore技术理解(上)

3.2 自动并行

随着深度学习的发展,为了实现更高的准确率和更丰富的应用场景,训练数据集和深度神经网络模型的规模日益增大。特别是自然语言处理(Natural Language Processing,NLP)领域,数据集的范围从 200MB 到 541TB 不等。模型的尺寸也从 BERT[18]的 3.4 亿个参数,Transformer-x l[19]的 8 亿 个参数,GPT-2 [19]的 150 亿个参数,到最近 NVIDIA Megatron-LM[20]的 80 多亿个参数。因此,在大型数据集上训练大型模型,不仅需要深度学习框架支持数据并行和模型并行,还需要支持混合并行。

当前的主流框架(如 TensorFlow[5]、Caffe[10]和 MXNet[7])大多通过手动切分深度神经网络模型来实现模型并行,但手动切分模型难度非常大,对开发者的要求非常高,需要有丰富经验的专家来操作。

实现混合并行(数据并行和模型并行同时进行)又极大增加了开发的复杂度。最近的研究成果[21–25] 提出了简化混合并行的方法,但这些方法在几个方面都存在局限性。

首先,这些方法在整个模型中,固定了对张量维数的切分策略,这可能导致切分策略是次优的,因为模型的不同部分适用不同的策 略。

第二,[22,24]不适合用于许多用于语言建模和行人重新识别(Re-identification,ReID)的深度神经网络,这些网络往往是非线性网络。

第三,[21]将划分策略搜索问题表述为混合整数程序,并使用现有的求解程序来求解,当处理大型模型时,速度会变得非常慢。此外,文献[21,24,25]旨在仅优 化通信或内存开销,可能不会减少训练时间。

MindSpore 的目标是在模型训练过程中允许并行转换。为此,在并行化策略搜索中引入了张量重排 布(Tensor Redistribution,TR),这使输出张量的设备布局在输入到后续算子之前能够被转换,如 3.2 图 3 中红色矩形所示。但是,对于复杂大型模型的搜索并行策略,在考虑张量重排布时,需要克服的挑战主要有两个。

首先,既然张量重排布将通信算子(例如 AllGather)引入数据流图,那么如何像普通算子一样自动地对通信算子求导呢?对于每个相应的前向算子,都需要获取反向算子,用来更新可训练参数。目前的框架需要专家在反向阶段手动添加 SEND 和 RECV 源语来传递梯度, 这对模型开发者来说是一项具有挑战性的工作,尤其是在模型比较复杂的情况下。

其次,随着张量重排布对策略空间的极大扩展,如何为复杂的大型模型高效地找到一个好的策略?在功能和效率方面,该算法需要为具有非线性结构的大模型快速找到一种策略。在性能方面,算法返回的策略应该会缩短端到端的训练时间。需要对运行成本进行仔细建模,这一过程也增加了人工成本。

MindSpore 针对上述两个挑战,推出了解决方案。为了实现通信算子的自动微分,MindSpore 定义 了通信算子的反向算子。例如,AllGatheris的反向算子为ReduceScatter,SEND的反向算子为RECV 后跟一个 ADD。定义这些反向算子十分重要,因为 Auto-diff 过程可以一次性地区分整个前向图,而无须跳过任何算子,这也是为什么 Auto-diff 是 Auto-parallel 后面一步的原因。

针对第二个挑战,在 同时考虑计算和通信开销的情况下,建立一个代价模型来选择一个好策略。为了快速地为复杂大图找到一个好策略,提出了几种方法:

一种是支持多图操作的算法,将原始图转换成线性图;

一种是策略分离机制,在保证返回解的精度的同时,有效地缩小搜索空间。例如,ResNet50 在 8 台设备上搜索并行策略的时间在 1s 内,而返回的解决方案确实缩短了训练时间。例如,当模型较大(类的 数量超过 128K)时,返回的解决方案与原始数据并行策略相比减少了大约 55%的训练时间。

代码 2 半自动并行配置中的并行过渡

1 class Submodel(nn.Cell):

2 def _init_(self, shape):

3 self.bn = BatchNorm(set_strategy={[4, 1]})

4 self.matmul = MatMul(set_strategy={[1, 1], [1, 4]})

5 self.W = Parameter(Tensor(shape), require_grad=True)

6 def construct(self, X):

7 Y = self.bn(X)

8 Z = self.matmul(y, self.W)

9 return Z

MindSpore技术理解(上)

MindSpore 较灵活,它支持用户指定的高级策略配置,称之为半自动并行(semi-auto-parallel)。在 代码 2 和图 3 中,展示了一个从数据到模型的并行转换的例子。该子模型的结构为 BatchNorm 算子后跟一个 MatMul 算子,广泛应用于 ResNet、ReID 等分类任务。在 BatchNorm 算子中,X 按行拆 分为四部分,数据可以并行,效率非常高。在 MatMul 算子中,可学习参数的权重 W 被分成四部分,模型可以并行,由于参数数量较多,这部分的模型并行更有效。由于 BatchNormi 的输出布局与 MatMul 的输入布局不同,所以框架插入了一个张量重排布(该例中为 AllGather 和 ConCat),这一过程对用户是透明的。用户也不必关注哪个设备运行了模型的哪个部分,框架会自动安排。

然而,不同的模型结构在每个算子中具有不同大小的参数,如图 4 所示,并且它们使用于不同的切分策略。在图 4(3)中,将第一算子配置为模型并行,将后续算子配置为数据并行,也需要插入张量重排布,这样可以获得更好的性能。

MindSpore技术理解(上)

在训练新模型时,多次配置 set_strategy,耗时耗力。在这种情况下,如果配置了自动并行 (auto-parallel),则不需要指定 set_strategy,该算法将找到一个有效的策略。例如,当 ResNet 中 的分类数量超过 130K 时,算法返回的策略导致在 50ms 内训练一个迭代。相比之下,原始数据并 行训练一次迭代超过 111ms。更详细的评估,见 8.1 节。

3.3 动态图

由于编译器能获得静态图的全局信息,所以静态图在大多数情况下都表现出更好的运行性能。而动态图可以保证更好的易用性,使用户能够更加方便地构建和修改模型。为了同时支持静态图和动态图,大多数先进的训练框架,需要维护两种自动微分机制,即基于 Tape 的自动微分机制和基于图的自动微分机制。从开发者的角度来看,维护两套自动微分机制成本较高;从用户的角度来看,在静态模式和动态模式之间切换也相当复杂。

MindSpore 采用了基于源码转换的自动微分机制,同时支持静态图和动态图,高效易用。在 MindSpore 中,称动态图为 Pynative 模式,因为代码使用 Python 解释器在这种模式下运行。如代 码 3 所示,从静态图模式切换到动态图模式只需要一行代码,反之亦然。此外,为了提高 Pynative 模式的运行效率,MindSpore 支持 staging 机制,如代码第 4 行所示。在函数(fc_relu)前添加 ms_function 装饰器,该函数将以静态图模式编译和运行。

代码 3 LeNet5 的 MindSpore 实现

1 from mindspore import context

2 import numpy as np

3 class LeNet5(nn.Cell):

4 @ms_function

5 def fc_relu(self, x):

6 x = self.relu(self.fc2(x))

7 x = self.fc3(x)

8 def construct(self, x):

9 x = self.max_pool2d(self.relu(self.conv1(x)))

10 x = self.max_pool2d(self.relu(self.conv2(x)))

11 x = self.flatten(x)

12 x = self.relu(self.fc1(x))

13 x = fc_relu(x)

14 return x

15 data=np.ones((batch_size,3,224,224),np.float32)*0.01

16 net = LeNet5()

17 # switch to Pynative mode

18 context.set_context(mode=context.PYNATIVE_MODE)

19 pynative_out = net(data)

20 # switch back to static graph mode

21 context.set_context(mode=context.GRAPH_MODE)

22 graph_out = net(data)

静态图和动态图的架构如 3.3 图 5 所示。在动态图模式中,框架遍历模型的所有算子,为每个算子 生成计算图,并将计算图传递给 GE 进行前向计算。在完成前向计算后,框架为模型生成前向和反向计算图,并将它们发送到 GE 中执行。由于 MindSpore 使用基于源码转换的自动微分机制,因此,在生成反向图时,可以省略用户为检查其模型而设置的所有代码,如 pdb 和 print。

MindSpore技术理解(上)