写点什么

AI for Science 系列(三):赛桨 PaddleScience 底层核心框架技术创新详解

  • 2023-02-22
    北京
  • 本文字数:2561 字

    阅读完需:约 8 分钟

框架技术创新

飞桨科学计算套件赛桨 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 个小算子。



图 左:CINN 架构图 中:gemm rewriter 示例 右:5 个小算子融合示例


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



(性能测试硬件平台:V100 16GB 单卡;Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz 40-core)


从实验结果可以看出,受益于新自动微分机制上的革新和自研编译器 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****

用户头像

还未添加个人签名 2022-12-26 加入

还未添加个人简介

评论

发布
暂无评论
AI for Science系列(三):赛桨PaddleScience底层核心框架技术创新详解_人工智能_飞桨PaddlePaddle_InfoQ写作社区