AI for Science系列（三）：赛桨PaddleScience底层核心框架技术创新详解

继上一篇典型案例及API功能介绍，本篇重点讲解飞桨核心框架为支持科学计算任务在技术上的创新工作与成果。

框架技术创新

飞桨科学计算套件赛桨PaddleScience底层技术依赖飞桨核心框架。为了有力地支撑科学计算任务高效训练与推理，飞桨核心框架在自动微分、编译器、执行器和分布式等多方面分别进行了技术创新。值得一提的是，自动微分在机制上的革新带来了框架计算量方面的优化，结合自研编译器CINN的加持，在典型科学计算任务2D定常Laplace模型上达到业界性能最优。

自动微分
飞桨核心框架在2.3版本中针对自动微分机制和自动微分API提供多种技术创新。

自动微分机制

自动微分在机制方面，通过设计自动微分基础算子体系，并在其上定义linearize和transpose程序变换规则。飞桨在静态图中新增对前向微分的支持，实现前反向不限阶的自动微分机制。创新后的自动微分机制具备良好的算子支持扩展性，目前已经支持全连接网络，并在2D定常Laplace、3D定常圆柱绕流等典型科学计算任务上完成验证工作。

目前的自动微分机制设计流程为：用户完成前向组网后，会首先将前向过程对应的网络转化到自动微分基础算子体系上，然后多次调用linearize和transpose两种程序变换实现高阶前向/反向微分，最后将网络整体转化到飞桨原生算子体系上完成组网或者对接编译器进行更多优化。该机制将计算拆分到更细的粒度，由于引入并合理使用切空间（前向微分空间）信息，相对于原有的自动微分机制有计算量上的优化，在2D定常Laplace模型上实测有最大1.58倍性能提升。

自动微分API

在自动微分API方面，飞桨为外部用户新增2个自动微分API，分别是前向自动微分forward_grad与反向自动微分grad两种接口，均支持高阶微分。为了方便用户使用，同时提供Jacobian、Hessian的求解，支持按行延迟计算，在复杂偏微分方程组中显著提升计算性能。另外，也公开了实现此次创新后的自动微分核心机制API jvp与vjp，当前所有API均放在incubate空间下，处于实验特性，会根据用户需求及具体应用场景持续演进。上述API列表以及计算示例如下表所示。

分布式自动并行

为了支持用户使用大规模网格点数据进行高效训练，PaddleScience采用高可扩展的飞桨自动并行技术，它能根据用户输入的串行网络模型和所提供集群资源信息自动进行分布式训练。基于统一分布式计算图和统一资源图设计，飞桨自动并行可支持任意并行策略和各类硬件集群资源上的分布式训练，能将原用户原计算图中每个串行tensor和串行operator自动切分成分布式tensor和分布式operator，并插入合适高效通信来保证与原串行计算一致，如下图所示。

目前飞桨自动并行支持半自动与全自动两种模式：

模式一：半自动模式下用户可以根据自己需要指定某些关键tensor和operator的切分方式，而框架会通过数据流分析方式进行自动推导剩余的tensor和operator；
模式二：全自动模式下不需要用户标记任何切分方式，所有tensor和operator都由框架基于Cost Model自适应选择最优的切分策略。基于自动并行所提供的数据并行、梯度累积和数据集并行加载等分布式能力，在3D圆柱绕流任务上32卡扩展效率可达到93%。结合自动并行，PaddleScience未来将更容易获得复杂并行策略支持。

编译器

自动微分机制将科学计算模型中的深度学习算子拆分成若干细粒度的基础算子，若直接计算，这些细粒度算子将使得模型的大部分时间被浪费在数据交换和执行调度而非实际计算上。这意味着我们并未完全利用到硬件的计算资源，因此模型的计算速度还仍有较大的提升空间。神经网络编译器在此类问题上有着显著的优势，包括针对计算图的通用Pass优化、算子融合以及自动代码生成等。飞桨神经网络编译器CINN即是解决此类问题的有力工具，它在消除大量无用数据交换和执行调度开销的同时，极大地降低算子开发所需的人力成本，相比于未开启编译加速的执行方式，开启CINN在2D定常Laplace模型上实测有最大3.72倍性能提升。

CINN为科学计算模型提供了一套完善的优化机制，典型的常量折叠，冗余算子消除，特定硬件的gemm rewriter库算子替换等优化Pass在CINN中均已实现。自动算子融合能力也在日趋强化，为科学计算等领域模型的训练性能提供更多的加速契机。以2D定常Laplace模型为例，最大的融合算子集合可包含94个小算子。

我们以2D定常Laplace模型作为统一测试模型，设定采样点数100*100，选取多组网络层数和隐藏层节点数的全连接网络训练2000 epoch，分别在JAX、飞桨静态图、飞桨静态图+新自动微分、飞桨静态图+新自动微分+CINN编译器几种模式下进行测试。以飞桨静态图模式作为baseline，计算其他模式相对飞桨静态图的加速比。

从实验结果可以看出，受益于新自动微分机制上的革新和自研编译器CINN的性能优化加持，在典型科学计算任务2D定常Laplace模型上达到业界性能最优。

执行器

飞桨中有两种数据结构可以表示网络模型，一种是Program，将模型表示成OP序列，方便在单机单卡下顺序执行OP；另一种是Graph，将模型表示成计算图结构，方便在多卡下并行调度OP，并默认开启许多图优化的IR pass以提升训练性能。由于Program相比Graph修改简单，可以方便地插入分布式需要的数据通信等OP，因而飞桨中许多重要的分布式功能都是基于Program结构进行开发。在科学计算项目中，我们对上述典型的科学计算模型实现Program和Graph两种表示结构的无损互转，使得计算图优化后的Graph可以转换成Program进行分布式训练，打通计算图优化和分布式训练结合的技术路径，充分利用两者的优化能力，大幅提升模型训练性能。

此外，飞桨框架在2.3版本中发布了全新的静态图执行器，其中实现高效的多stream和多线程异步调度组件，使得模型训练更加性能优越和易扩展，并已在单机单卡场景下默认使用。针对PaddleScience，我们通过对静态图新执行器在对接Graph IR体系和支持分布式训练等方面的功能适配和扩充，使图优化后的模型可以利用新执行器进行分布式训练。

拓展阅读
AI+Science系列（一） ：飞桨加速CFD（计算流体力学）原理与实践
AI+Science系列（二）：国内首个基于AI框架的CFD工具组件！赛桨v1.0 Beta API介绍以及典型案例分享！
赛桨PaddleScience v1.0 Beta：基于飞桨核心框架的科学计算通用求解器

相关地址
飞桨AI for Science共创计划
https://www.paddlepaddle.org.cn/science
飞桨PPSIG-Science小组
https://www.paddlepaddle.org.cn/specialgroupdetail?id=9