Posted in

灵活可扩展的新一代 IR 技术——飞桨框架 3.0 基石技术深度解析_AI阅读总结 — 包阅AI

包阅导读总结

思维导图:

文章地址:https://mp.weixin.qq.com/s/OkVVsPSlyRkJz6aXzPIMug

文章来源:mp.weixin.qq.com

作者:飞桨

发布时间:2024/7/31 12:34

语言:中文

总字数:5551字

预计阅读时间:23分钟

评分:82分

标签:深度学习框架,飞桨框架3.0,中间表示(IR),性能优化,大模型训练


以下为原文内容

本内容来源于用户推荐转载,旨在分享知识与观点,如有侵权请联系删除 联系邮箱 media@ilingban.com

为了让优秀的飞桨开发者们掌握第一手技术动态、让企业落地更加高效,飞桨官方在7月至10月特设《飞桨框架3.0全面解析》系列技术稿件及直播课程,邀请百度飞桨核心团队数十位工程师齐上阵,技术解析加代码实战,带大家掌握包括核心框架、分布式计算、产业级大模型套件及低代码工具、前沿科学计算技术案例等多个方面的框架技术及大模型训推优化经验。

本文是该系列第二篇技术稿件,文末附对应直播课程详情。

近些年来,越来越多的框架和研究者将编译器技术引入到深度学习的神经网络模型优化中,并在此基础上借助编译器的理念、技术和工具对神经网络进行自动优化和代码生成。作为深度学习框架性能优化、推理部署、编译器等方向的重要基石——计算图中间表示(Intermediate Representation,即 IR),在大模型场景下对其灵活性、扩展性、完备性也提出了更高的要求,即如何便捷地支撑大模型自动并行下丰富的切分策略,如何更低成本地实现性能优化 Pass 等。

在3.0版本下,飞桨充分考虑大模型对 IR 的需求,在基础架构层面规范了中间表示 IR 定义,全架构统一表示,实现了推理、分布式、编译器等上下游各方向共享开发成果。飞桨的新一代 IR 架构通过更加完备且鲁棒的语义表达能力、训推全架构统一表示和高效可插拔的性能优化策略(Pass)开发机制,无缝对接神经网络编译器实现自动性能优化和多硬件适配,加速大模型训推流程

对于一个特定的深度学习任务模型,从算法工程师的模型视角来看,主要涉及三个核心概念:数据集(Dataset)、网络模型(Model)、迭代训练(Training);但从研发工程师的框架视角来看,主要有两个核心概念:张量(Tensor)和算子(Operator)。广义上,算子是“计算操作”的一种表示,比如矩阵乘算子 Matmul,有两个输入和一个输出。狭义上,算子包含计算表示和 Kernel 定义,其中 Kernel 与硬件有很强的关联性,同一个算子表示可以对应多个 Kernel 定义。比如 Matmul 算子有 CPU Kernel,也有 GPU、XPU、NPU 等硬件上 Kernel。

飞桨采用了动静统一的技术架构设计,以动态图为默认编程范式。其即时执行的特点给开发者提供了 Pythonic 的流畅使用体验,同时提供了动转静功能(如飞桨的@to_static)支持用户可以一键转为静态图执行。用户的模型代码在静态图模式下会先记录所有的张量和算子操作,构建一个完整的计算图(即 IR 中间表示),然后借助框架的内核执行器在给定的硬件设备上执行训练或者推理。

在深度学习框架中,中间表示在架构层面承担着非常重要的角色,主要包括:

  • 推理部署:模型训练仅是深度学习链条中的一环,如何结合业务落地,导出模型推理部署开发给用户使用,与计算图中间表示息息相关。

  • 性能优化:IR 中包含了计算图的全局信息(如上下游算子的邻接关系等),更有利于进行图优化,如常量折叠、算子融合、Inplace 策略等。

  • 神经网络编译器:将 IR 中的全局信息传递给飞桨 CINN 编译器,可实现子图自动融合、Kernel 自动生成,节省工程师 80% 的时间达到人工调优相似的性能。

  • 社区生态:IR 表示本质是一层“协议”的概念,借助飞桨套件工具(如X2Paddle、Paddle2ONNX)与其他开源生态自由转换。

飞桨新的一代 IR 是基础架构层面的升级,对于用户在 API 层面的使用是无感的,用户可保持之前动转静(即 paddle.jit.to_static)或静态图代码不变,在3.0-Beta 下仅需额外通过export FLAGS_enable_pir_api=1开启新 IR 功能即可,如下是一个简单的使用样例。

import unittestimport numpy as npimport paddlefrom paddle.static import InputSpecclass SimpleNet(paddle.nn.Layer):    def __init__(self):        super().__init__()    def forward(self, x, y):        z = x + y        out = paddle.nn.functional.relu(z)        return outspecs = [InputSpec(shape=(-1, -1)), InputSpec(shape=(-1, -1))]net = paddle.jit.to_static(SimpleNet(), specs)x = paddle.rand(shape=[16, 64], dtype=paddle.float32)y = paddle.rand(shape=[16, 64], dtype=paddle.float32)out = net(x, y)print(out)

将上述文件保存为 test_add_relu.py,执行如下命令:FLAGS_enable_pir_api=1 python test_add_relu.py即可。开发者可额外指定 GLOG_v=6输出日志,查看新一代 IR 下的 Program 表示,如下所示,在动转静或静态图模式下,用户的代码经过组网 API 下会先生成 Operator Dialect 下计算图表示,在执行时飞桨会将其转换为给定硬件下的 Kernel Dialect,然后交给执行器去依次调度对应的 PHI 算子库,计算最终结果。

{     (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[-1,-1],stop_gradient:[true]} : () -> builtin.tensor<-1x-1xf32>    (%1) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"y",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[-1,-1],stop_gradient:[true]} : () -> builtin.tensor<-1x-1xf32>    (%2) = "pd_op.add" (%0, %1) {stop_gradient:[true]} : (builtin.tensor<-1x-1xf32>, builtin.tensor<-1x-1xf32>) -> builtin.tensor<-1x-1xf32>    (%3) = "pd_op.relu" (%2) {stop_gradient:[true]} : (builtin.tensor<-1x-1xf32>) -> builtin.tensor<-1x-1xf32>    () = "builtin.shadow_output" (%3) {output_name:"output_0"} : (builtin.tensor<-1x-1xf32>) ->}
{ (%0) = "data(phi_kernel)" () {dtype:(pd_op.DataType)float32,kernel_key:<backend:Undefined|layout:Undefined(AnyLayout)|dtype:float32>,kernel_name:"data",name:"x",op_name:"pd_op.data",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[-1,-1],stop_gradient:[true]} : () -> undefined_tensor<-1x-1xf32> (%1) = "shadow_feed(phi_kernel)" (%0) {kernel_key:<backend:GPU|layout:Undefined(AnyLayout)|dtype:float32>,kernel_name:"shadow_feed",op_name:"pd_op.shadow_feed"} : (undefined_tensor<-1x-1xf32>) -> gpu_tensor<-1x-1xf32> (%2) = "data(phi_kernel)" () {dtype:(pd_op.DataType)float32,kernel_key:<backend:Undefined|layout:Undefined(AnyLayout)|dtype:float32>,kernel_name:"data",name:"y",op_name:"pd_op.data",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[-1,-1],stop_gradient:[true]} : () -> undefined_tensor<-1x-1xf32> (%3) = "shadow_feed(phi_kernel)" (%2) {kernel_key:<backend:GPU|layout:Undefined(AnyLayout)|dtype:float32>,kernel_name:"shadow_feed",op_name:"pd_op.shadow_feed"} : (undefined_tensor<-1x-1xf32>) -> gpu_tensor<-1x-1xf32> (%4) = "add(phi_kernel)" (%1, %3) {kernel_key:<backend:GPU|layout:NCHW|dtype:float32>,kernel_name:"add",op_name:"pd_op.add",stop_gradient:[true]} : (gpu_tensor<-1x-1xf32>, gpu_tensor<-1x-1xf32>) -> gpu_tensor<-1x-1xf32> (%5) = "relu(phi_kernel)" (%4) {kernel_key:<backend:GPU|layout:NCHW|dtype:float32>,kernel_name:"relu",op_name:"pd_op.relu",stop_gradient:[true]} : (gpu_tensor<-1x-1xf32>) -> gpu_tensor<-1x-1xf32> () = "builtin.shadow_output" (%5) {output_name:"output_0"} : (gpu_tensor<-1x-1xf32>) ->}

在深度学习框架 IR 概念中,“顺序性”和“图语义”是两个非常高频常用的概念。旧的中间表示体系由“顺序性”ProgramDesc 和“图语义”Graph 两个核心类共同承载。用户在静态图 API 或者动转静模块下,产生的中间表示是 Op-by-Op 的 Program,如果要应用更高层面的优化策略(比如算子融合、inplace 策略、剪枝等),框架会将由 Program 构造出 Graph,其由数据节点、算子节点和彼此关联的边构成。

在新的 Paddle IR 中,飞桨在底层抽象了一套高度可扩展的基础组件,包括 Type、Attrbute、Op、Trait 和 Interface,并引入了 Dialect 的概念,支持开发者灵活扩展、自由定制,提供了完备鲁邦的语义表达能力;在模型表示层,通过多 Dialect 模块化管理,统一多端表示,实现了训推一体的全架构统一表示,无缝衔接组合算子和编译器,支持自动优化和多硬件适配;在图变换层,通过统一底层模块,简化基础概念,向用户提供了低成本开发、易用高性能、丰富可插拔的 Pass 优化机制。飞桨的新一代的 IR 表示坚持 SSA(静态单赋值)原则,模型等价于一个有向无环图。并以 Value、Operation 对计算图进行抽象:

  • Operation 表示计算图中的节点:一个 Operation 表示一个算子,它里面包含了零个或多个 Region;Region 表示一个闭包,它里面包含了零个或多个 Block;Block 表示一个符合 SSA 的基本块,里面包含了零个或多个 Operation;三者循环嵌套,可以实现任意复杂的语法结构。

  • Value 表示计算图中的有向边:用来将两个 Operaton 关联起来,描述了程序中的 UD 链(即 Use-Define 链);OpResult 表示定义端,定义了一个 Value,OpOperand 表示使用端,描述了对一个 Value 的使用。

如上左图所示,新一代 IR 的整体设计自底向上分为三层:

飞桨提供了 Trait 和 Interface 两种重要机制实现了对算子 Op 的特征和接口的抽象标记。比如 InplaceTrait 表示一个 Op 具有 Inplace 特征, InferShapeInterface 表示一个算子定义了 InferShape 函数接口等,这二者都是可以任意扩展的,仅派生自相应的基类、遵循相应的实现规则即可;并对算子体系下核心概念抽象出 Type、Attrbute、Op,这三者是基于 Trait 和 Interface 进行定义的。Dialect 用来对 Type、Attribtue、Op 做模块化管理, 比如 BuiltinDialect、DistDialect、CinnDialect 等等。一个 Dialect 里面包含了一系列的 Type、Attribtue、Op 的定义。相应的,每个 Type、Attribtue、Op 都是定义在某个唯一的 Dialect 里面。对整个 IR 框架而言, Dialect 是可以随意插拔的,也是可以任意扩展的。

这一层是 IR 适应多种场景的基础。这一层的每一个要素都是可定制化扩展,一般情况下,针对一个具体的场景(比如分布式、编译器)只需要在其 Dialect 中定义特定的 Trait、Interfce 以及 Type、Attribute、Op 即可。

飞桨通过不同层级的 Dialect 来管理框架内不同领域的算子体系,比如 Built-in 下的 Shape Dialect 和 Control Flow Dialect(分别用于形状符号推导和控制流表示)、与 PHI 算子库执行体系相关的 Operator Dialect 和 Kernel Dialect、与神经网络编译器领域相关的 CINN Dialect 等。在飞桨神经网络编译器中,主要以计算图 Operator Dialect 为输入,经过组合算子和 Pass Pipline 后,会转换为 CINN Dialect,并附加 Shape Dialect 中的符号信息,最后会 Lowering 成编译器的 AST IR。

上述这些多层级的 Dialect 内的算子 Op 会组成 Program ,并用来表示一个具体的模型,它包含两部分:计算图和权重

  • Value、Operation 用来对计算图进行抽象。Value 表示计算图中的有向边,用来将两个 Operaton 关联起来,描述了程序中的 UD 链 ,Operation 表示计算图中的节点。一个 Operation 表示一个算子,它里面包含了零个或多个 Region 。Region 表示一个闭包,它里面包含了零个或多个 Block。Block 表示一个符合 SSA 的基本块,里面包含了零个或多个 Operation 。三者循环嵌套,可以实现任意复杂的语法结构。

  • Weight 用来对模型的权重参数进行单独存储,这也是深度学习框架和传统编译器不一样的地方。传统编译器会将数据段内嵌到程序里面。这是因为传统编译器里面,数据和代码是强绑定的,不可分割。但是对神经网络而言,一个计算图的每个 epoch 都会存在一份权重参数,多个计算图也有可能共同一份权重参数。

Pass 的核心是子图匹配和替换(即图变换),是将一个 Program 通过某种规则转换为另一个新的 Program。IR 中包含了计算图中全局信息,如上下游算子的邻接关系等,更有利于进行图优化,比如常量折叠、算子融合,Inplace 策略等。我们将在第三章节详细介绍 Pass 开发机制。

飞桨以 Paddle IR 为基础,提供了高效可插拔的性能优化 Pass 策略开发机制,在训练、推理、分布式、编译器等多领域场景下可复用,且友好支持多硬件定制化。在新 IR 体系下,飞桨支持2种 Pass 开发范式:Pattern Rewriter 和 DRR,充分兼顾自定义灵活性和开发易用性,实现了 Pass 开发成本降低58%;应用于推理场景,超过84%的模型推理加速超10%

Pass 的核心是一个“子图匹配和替换”的过程,一般需要开发者定义待替换子图的源 Pattern,以及期望替换的目标 Pattern。在 Pattern Rewriter 机制下,分别对应于实现 Match 和 Rewrite 逻辑。

举个例子,假设在 Kernel Dialect 里,我们期望移除冗余的 Feed 算子,以提升执行效率。则需要如下两个步骤。

1.继承 PatternRewritePass类:只需要在 InitializePatterns 里添加目标算子的 Pattern 即可,开发者无需关心计算图遍历实现;

2.定义Match 和 Rewrite逻辑:新增目标算子的Pattern,实现 Match 和 Rewrite 接口函数。

Declarative Rewrite Rule 机制(即 DRR)是一种三段式 Pass 开发范式,让开发者更聚焦于 Pass 逻辑的处理,无需关心底层 IR 的实现细节,所见即所得

举个例子,在混合精度训练的计算图中,常会出现连续的 DataType 相关的 Cast 操作。从优化角度,类似连续的 Cast 操作完全可以简化成一个 Cast 即可,如下图左侧示意图。

基于 DRR 的 Pass 机制,我们只需要如下两个简单的步骤

1.继承 DrrPatternBase 模板类:此模版类提供了一系列的基础方法,帮助用户用更简洁的形式描述中间 DAG 子图表示的 Pattern 信息。

2.定义Source/Target Pattern:通过 DrrPatternContext来创建、描述源、目标Pattern,仅需要聚焦算子的输入、输出关系即可。

什么是 Dialect 呢?可以简单理解为“域”的概念,是高度灵活性的核心支撑。飞桨新 IR 下可以定义不同的 Dialect,如用于编译期算子定义的 Operator Dialect、用于执行期的 Kernel Dialect、用于 AI 编译器的 CINN Dialect 等。

在飞桨神经网络编译器中,以计算图 Operator Dialect 为输入,经过组合算子和 Pass Pipline 后,会转换为 CINN Dialect,并附加 Shape Dialect 中的符号信息,最后会 Lowering 编译器 AST IR。各阶段涉及的 Dialect 主要为:

  • 计算图:为 Operator Dialect,是原始计算图表示,包含未拆解的大算子,以及原始的动态 Shape,无符号约束信息;

  • 符号推导:为 Shape Dialect,经过组合算子拆分后,会调用基础算子的 InferSymbolicShape 逻辑,记录张量 Tensor 的符号信息和逻辑计算表达式;

  • 编译器高层 IR:为 CINN Dialect,编译器在算子层面有特殊的定义,比如“属性信息静态化”,以满足更高的优化潜力,框架会使用 Dialect 之间的 Convert Pass 将 Operator Dialect 中的特定算子转换为 CINN Dialect 中对应的算子;

在 3.0 框架中,飞桨对150+个基础算子和典型算子新增支持了动态 Shape 的推导逻辑(即 InferSymbolicShape),可动态地随计算图 Pass 的应用来自适应且 Lazy 地更新计算图中间表示的符号信息。

飞桨 3.0 Beta 版本下,推出了新一代中间表示 PIR,这项技术对底层的核心概念如 Operation、Attribute 等进行了系统性的抽象,为开发者提供了灵活的基础组件;同时,通过引入 Dialect 这一概念,飞桨能够全面、分层次管理框架各模块对中间表示的需求,并支持开发者根据需求定制化扩展 Dialect,显著提升了框架的扩展性。PIR 遵循 SSA(即 Static Single Assignment)原则,统一了顶层结构,实现“算子顺序性”和“计算图语义”的兼容表示。此外,PIR 还提供了更加简洁、低成本的 Pass 开发流程(DRR 和 pattern write),并内置了一系列丰富且功能完备的 Pass 优化策略,实现了 Pass 开发成本降低58%;应用于推理场景,超过84%的模型推理加速超10%,为大模型的极致性能优化提供了强有力支撑。

7月-10月特设《飞桨框架3.0全面解析》直播课程,邀请百度飞桨核心团队数十位工程师倾囊相授,技术解析+代码实战,带大家掌握核心框架、分布式计算、产业级大模型套件及低代码工具、前沿科学计算技术案例等多个方面的框架技术及大模型训推优化经验,实打实地帮助大家用飞桨框架3.0在实际开发工作中提效创新!

为了让优秀的飞桨开发者们掌握第一手技术动态、让企业落地更加高效,根据大家的呼声安排史上最强飞桨技术大餐!涵盖飞桨框架3.0、低代码开发工具 PaddleX、大语言模型开发套件 PaddleNLP、多模态大模型开发套件 PaddleMIX、典型产业场景下硬件适配技术等多个方向,一起来看吧!

温馨提示:以上仅为当前筹备中的部分课程,

如有变动,敬请谅解。

【前序技术稿件】

飞桨新一代框架3.0:“动静统一自动并行、大模型训推一体”等新特性构筑大模型时代核心生产力

【3.0-Beta 视频教程】

https://aistudio.baidu.com/course/introduce/31815

【3.0-Beta 官方文档】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/index_cn.html

【开始使用】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/overview_cn.html#jiukaishishiyong

【动转静 SOT 原理及使用】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/sot_cn.html

【自动并行训练】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/auto_parallel_cn.html

【神经网络编译器】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/cinn_cn.html

【高阶自动微分功能】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/higher_order_ad_cn.html

【PIR 基本概念和开发】

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/paddle_ir_cn.html

【飞桨官网】

https://www.paddlepaddle.org.cn/

【企业合作入口】

https://paddle.wjx.cn/vm/m3sxpfF.aspx#