灵活可扩展的新一代IR技术——飞桨框架3.0基石技术深度解析

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

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

前言

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

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

新一代 IR 设计

什么是中间表示 IR?

手写字体识别任务模型
算子表示
对于一个特定的深度学习任务模型,从算法工程师的模型视角来看,主要涉及三个核心概念:数据集(Dataset)、网络模型(Model)、迭代训练(Training);但从研发工程师的框架视角来看,主要有两个核心概念:张量(Tensor)和算子(Operator)。广义上,算子是「计算操作」的一种表示,比如矩阵乘算子Matmul,有两个输入和一个输出。狭义上,算子包含计算表示和 Kernel 定义,其中 Kernel 与硬件有很强的关联性,同一个算子表示可以对应多个Kernel定义。比如 Matmul 算子有 CPU Kernel,也有 GPU、XPU、NPU等硬件上 Kernel。
飞桨采用了动静统一的技术架构设计,以动态图为默认编程范式。其即时执行的特点给开发者提供了 Pythonic 的流畅使用体验,同时提供了动转静功能(如飞桨的@to_static)支持用户可以一键转为静态图执行。用户的模型代码在静态图模式下会先记录所有的张量和算子操作,构建一个完整的计算图(即 IR 中间表示),然后借助框架的内核执行器在给定的硬件设备上执行训练或者推理。
静态图模式下执行流程

为什么需要 IR ?

在深度学习框架中,中间表示在架构层面承担着非常重要的角色,主要包括:
推理部署:模型训练仅是深度学习链条中的一环,如何结合业务落地,导出模型推理部署开发给用户使用,与计算图中间表示息息相关。
性能优化:IR 中包含了计算图的全局信息(如上下游算子的邻接关系等),更有利于进行图优化,如常量折叠、算子融合、Inplace 策略等。
神经网络编译器:将 IR 中的全局信息传递给飞桨CINN编译器,可实现子图自动融合、Kernel 自动生成,节省工程师 80% 的时间达到人工调优相似的性能。
社区生态:IR 表示本质是一层「协议」的概念,借助飞桨套件工具(如X2Paddle、Paddle2ONNX)与其他开源生态自由转换。
性能优化Pass&社区生态

初窥 Paddle IR

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

# test_add_relu.py

import unittest
import numpy as np
import paddle
from paddle.static import InputSpec

class SimpleNet(paddle.nn.Layer):
    def __init__(self):
        super().__init__()


    def forward(self, x, y):
        z = x + y
        out = paddle.nn.functional.relu(z)
        return out
# Step 1: 构建模型对象,并应用动转静策略
specs = [InputSpec(shape=(-1, -1)), InputSpec(shape=(-1, -1))]
net = paddle.jit.to_static(SimpleNet(), specs)

# Step 2: 准备输入,执行 forward
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 算子库,计算最终结果。

{ // Operator Dialect
    (%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>) ->
}

// IR after lowering
{ // Kernel Dialect
    (%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>) ->
}

PIR 架构设计

在深度学习框架 IR 概念中,「顺序性」和「图语义」是两个非常高频常用的概念。旧的中间表示体系由「顺序性」ProgramDesc 和「图语义」Graph 两个核心类共同承载。用户在静态图 API 或者动转静模块下,产生的中间表示是 Op-by-Op 的 Program,如果要应用更高层面的优化策略(比如算子融合、inplace 策略、剪枝等),框架会将由 Program 构造出 Graph,其由数据节点、算子节点和彼此关联的边构成。
Paddle IR 分层架构图
在新的 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 的整体设计自底向上分为三层:
1.灵活的基础组件
飞桨提供了 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即可。
2.多层级的 Dialect
飞桨通过不同层级的 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 都会存在一份权重参数,多个计算图也有可能共同一份权重参数。
3.功能完善的 Pass 体系
Pass 的核心是子图匹配和替换(即图变换),是将一个 Program 通过某种规则转换为另一个新的 Program。IR 中包含了计算图中全局信息,如上下游算子的邻接关系等,更有利于进行图优化,比如常量折叠、算子融合,Inplace 策略等。我们将在第三章节详细介绍Pass开发机制。

Pass 体系介绍

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

Pattern Rewriter 机制

Pass 的核心是一个「子图匹配和替换」的过程,一般需要开发者定义待替换子图的源Pattern,以及期望替换的目标Pattern。在Pattern Rewriter 机制下,分别对应于实现 Match 和 Rewrite 逻辑。
举个例子,假设在Kernel Dialect里,我们期望移除冗余的 Feed 算子,以提升执行效率。则需要如下两个步骤:
在这里插入图片描述

继承 PatternRewritePass类:只需要在 InitializePatterns 里添加目标算子的 Pattern 即可,开发者无需关心计算图遍历实现;
定义Match 和 Rewrite逻辑:新增目标算子的Pattern,实现 Match 和 Rewrite 接口函数

DRR 机制

Declarative Rewrite Rule 机制(即DRR)是一种三段式 Pass 开发范式,让开发者更聚焦于Pass逻辑的处理,无需关心底层IR的实现细节,所见即所得。
举个例子,在混合精度训练的计算图中,常会出现连续的 DataType 相关的 Cas t操作。从优化角度,类似连续的Cast操作完全可以简化成一个 Cast 即可,如下图左侧示意图。
在这里插入图片描述
在这里插入图片描述
基于 DRR 的Pass机制,我们只需要如下两个简单的步骤
继承DrrPatternBase模板类:此模版类提供了一系列的基础方法,帮助用户用更简洁的形式描述中间DAG子图表示的Pattern信息
定义Source/Target Pattern:通过DrrPatternContext来创建、描述源、目标Pattern,仅需要聚焦算子的输入、输出关系即可

Dialect 最佳实践

什么是 Dialect 呢?可以简单理解为「域」的概念,是高度灵活性的核心支撑。飞桨新 IR 下可以定义不同的 Dialect,如用于编译期算子定义的 Operator Dialect、用于执行期的 Kernel Dialect、用于 AI 编译器的 CINN Dialect等。
“从类型的角度,恰当分层的软件栈需要支持对张量、buffer、向量、标量等进行建模,以及一步步分解和Lowering”
“从操作的角度,我们需要计算和控制流、控制流可以是显式的基础块跳转,也可以内含于结构化操作中”
编译器架构概览
在飞桨神经网络编译器中,以计算图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版本下,推出了新一代中间表示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#

  • 11
    点赞
  • 8
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值