分类

Posted in

统一多场景自动编译加速——支持动态 shape 场景，一套架构搞定训推需求_AI阅读总结 — 包阅AI

by baoyueai

包阅导读总结

1. 飞桨、编译器技术、模型优化、动态 shape、性能提升

2. 飞桨为让开发者掌握技术动态、企业落地更高效，推出系列技术稿件及直播课程。本文介绍飞桨神经网络编译器，其分为前端和后端，满足训推需求和动态 shape 场景，通过优化技术提升性能，还引入自动调优模块，能生成最优策略。

3.

– 飞桨框架 3.0 系列技术解读与直播课程

– 目的：让开发者掌握技术动态，企业落地更高效

– 内容：技术解析加代码实战，涵盖核心框架等多方面

– 深度学习依赖编译器技术的原因

– 硬件发展趋势：算力发展快，访存性能等影响大，手工融合覆盖窄，编译器自动融合优化可提升性能

– 模型发展趋势：模型结构多样，依赖编译器通用优化降低研发成本

– 多硬件优化：不同硬件特性和需求不同，编译器可减少开发工作量

– 飞桨神经网络编译器

– 整体架构：分为前端和后端

– 前端：基于 PIR 进行图级别优化，包括组合算子拆分、重计算策略、图优化 Pass、算子融合、维度推导等

– 后端：负责将前端处理后的 IR 转换为目标硬件可执行代码，包括基于硬件特性的 IR 优化、高效内存管理和代码生成等

– 自动调优模块：分析输入 shape，生成最优 Schedule 变换策略，Codegen 打印硬件指令并编译

思维导图：

文章地址：https://mp.weixin.qq.com/s/xVcHSym-fZvLeRM9AijNbw

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

作者：全栈布局的

发布时间：2024/8/7 12:23

语言：中文

总字数：6251字

预计阅读时间：26分钟

评分：88分

标签：深度学习,编译器技术,飞桨神经网络编译器(CINN),训练与推理优化,自动编译加速

以下为原文内容

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

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

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

近些年，在深度学习场景如何更好的借助编译器技术来提升任务表现变得越来越重要，在工业界和学术界均有大量基于编译器技术的探索和落地。为什么深度学习场景非常依赖编译器技术，整体趋势上的有3大原因：

1）硬件发展趋势：结合硬件发展历史和技术演进特点，算力发展速度远大于访存性能、CPU 性能和总线带宽；其中访存性能影响访存密集型算子(norm 类，activation 等)性能，CPU 性能和总线带宽影响调度性能。通过手工融合的方式，覆盖场景较窄，优化成本较高。基于编译器的自动融合的通用优化技术，将多个算子融合成一个大算子，通过减少访存量和算子数量，能够大幅提升模型端到端性能，编译器技术会成为深度学习框架标配组件。

2）模型发展趋势：模型结构存在多样性的特点，多样性的需求非常依赖编译器的通用优化来降低整体的研发成本。虽然在文本领域模型主流结构趋同，但是 Transformer 架构的推理成本较高，研究者也在积极探索如 Mamba 等结构，从模型结构角度来降低推理成本；在多模态、自动驾驶和科学计算等领域，模型结构还在探索中。

3）多硬件优化：当前市面存在有多款硬件，不同的硬件平台有不同的特性和优化需求，每个硬件均需要投入大量的人力进行优化，借助编译器技术，理论上仅需实现新硬件 IR 层面的对接，以及相应的硬件 IR 优化策略就能完成与深度学习框架的对接，相比于实现几百个硬件 Kernel，开发的工作量会大幅减少。

让我们通过一个实例来阐释针对模型多样性的趋势，自动融合优化的好处。我们以 Llama 模型中经常使用的 RMS Normalization （Root Mean Square Layer Normalization）为例，其计算公式相对简单明了。

假设我们需要是实现 RMS Normalization 的计算，最简单的办法是，我们可以使用飞桨框架提供的张量运算开发接口，调用平方、求和、除法、开根号等操作来完成，代码如下：

class RMSNorm(paddle.nn.Layer):    def __init__(self):        super().__init__()        self.variance_epsilon = 1e-6        self.size = 768        self.weight = paddle.create_parameter(            shape=[self.size],            dtype=paddle.get_default_dtype(),            default_initializer=nn.initializer.Constant(1.0),        )    def forward(self, x):        variance = x.pow(2).mean(-1, keepdim=True)        x = paddle.rsqrt(variance + self.variance_epsilon) * x        return x * self.weight

上述代码开发简单，但是性能较差，且显存占比较多；开发者可以实现一个 FusedRMSNorm 的实现，但是对于开发者要求更高，成本也更高。

借助神经网络编译器自动融合优化技术，我们能够在维持高度灵活性基础上，大幅降低了用户的优化成本，且实现性能的显著提升。以下 A100平台上 RMSNorm 算子的性能测试结果便是一个明证：相较于采用 Python 开发接口组合实现的方式，经过编译优化后的算子运行速度提升了4倍；即便与手动算子融合的方式相比，也实现了14%的性能提升。这一成果充分展示了飞桨框架在保证灵活性基础上，也低成本的取得了较好的性能表现。

神经网络编译器（也称为深度学习编译器）是一种专门为深度学习模型优化和部署而设计的工具，用于提高模型的计算效率、降低内存占用、加速训练推理过程等。其功能是将高层次的深度学习模型转换为低层次的、高效的、底层硬件可执行的代码。简单来说，深度学习编译器在深度学习框架和底层硬件之间充当了“翻译”的角色，能够将用户定义的神经网络模型描述转化为底层硬件能够理解和执行的指令。编译器在实现这种转换的过程中，应用了一系列优化技术，以提高模型在各种硬件平台上（包括 CPU、GPU 等硬件）的执行效率。

飞桨神经网络编译器（CINN, Compiler Infrastructure for Neural Networks）整体架构如上图所示，大体可以分为两大模块，分别是编译器前端和编译器后端。飞桨神经网络编译器充分考虑训练和推理场景的需求，一套架构同时满足训推需求。为了满足动态 shape 场景，设计了符号推导模块，高效处理动态 shape 问题。执行体系充分考虑和算子库、执行器高效整合，框架原生支持编译器体系，调度性能极致。

一般来说编译器前端核心功能并进行图级别的优化，CINN 作为飞桨框架原生的编译器，可以直接使用飞桨框架提供的模型加载和中间表示（Paddle IR，简称 PIR）组件，因此 CINN 前端的主要功能是基于 PIR 进行图层级别的优化，并对子图进行划分为后端高性能 Kernel 代码生成提供支持。CINN 前端关键的流程可分为三部分：

飞桨框架中将算子划分为基础算子（也称作原子算子，从性能和模型收敛正确性的角度，该类算子不适合进一步拆分）和非基础算子（非基础算子可以通过基础算子组合实现）两大类，由于非基础算子数量较多，会大幅增加编译器融合优化的难度，因此我们将组合算子拆分为等价的基础算子组合，借助编译器的自动融合技术可大幅提升性能的可优化空间。

经过组合算子拆分之后，在训练场景，由于有反向逻辑的存在，需要将前向算子的一些临时变量，保留到反向阶段，这会增加前向阶段 Kernel 输出和反向阶段 Kernel 输入的个数，通过使用重计算策略可以减少临时变量的输出和输入的个数，能够显著提升端到端的性能。

在计算图层级进行 PIR 的 Pass 优化，常见的图优化 Pass 包括：常量折叠、死代码消除（DCE）、公共子表达式消除（CSE）、冗余算子消除、算子计算合并等。

算子融合是编译器前端非常重要的一个功能，主要是将多个算子打包到一个子图中（对应为一个 FusionOp），交给编译器后端生成一个高效的硬件相关计算 Kernel。算子融合的本质是通过 IO 优化加速访存密集算子，如果我们将两个连续 Kernel 合并为一个 Kernel 调用，我们会减少中间变量的读写开销，因此在访存密集型的2个 Op 上，融合可以获取更高的性能。举个例子，如下图：

我们有两个算子 Relu 和 Scale，因为两个算子都是 IO 密集型算子（计算复杂度不高）。正常情况下我们需要读取 A 和 B 一次，写 B 和 C 一次。但是对于融合之后的 Kernel（右图）而言，我们只需要读取 A 和写 C 一次，这样我们通过算子融合可以取得更少的访存次数，在 IO 密集算子而言，可以极大提高性能。具体的算子融合策略实现非常复杂，这里不做展开介绍，感兴趣的读者可以阅读相关源码（地址请见文末）。

编译器期望通过尽可能大粒度的算子融合来保障端到端性能，但是大粒度融合强依赖正确性的保证。通过将当前可融合的算子分为两大类：TrivialOp和ReduceOp，通过理论论证 TrivialOp 和 ReduceOp 的融合正确性，来保障最终的正确性。

  • T+T 融合：使用 inline 代换可以实现 T * T 的融合

  • T+R 融合：使用 inline 代换来实现实现 T+R 的融合

  • R+R 融合：融合场景较复杂的一个部分，为了解决 inline 不封闭性质，无法通过维护一个 Op 来包含所有的 ReduceOp + ReduceOp 的信息，所以需要额外定义一个新的结构 ReduceTree 来实现封闭性，然后定义一个 Lower 操作在 ReduceTree 上生成一个融合的 Op Kernel 代码，就可以实现多个 Reduce 的融合

  • R+T 融合：是否需要融合需要考虑性能，因为 Reduce 的重计算开销太大，仅融合 T 的输出尺寸比 R 的输入尺寸小的

通过这种融合正确性的保障，我们可以在保障正确性的前提下，最大程度保障融合 kernel 性能，下面是我们在一些子图上，对比和 PyTorch 融合粒度的对比，CINN 能够融合为一个 Kernel。

维度推导是当网络中存在动态 shape 时，在之前的静态图模式下，常用-1来表示，但是这种-1形式，有效信息太少，无法有效支持后端做性能优化，因此需要有一个完整的维度推导机制，提供尽可能多的确定信息，来提升动态 shape 下 kernel 的性能。

对于维度推导，有两个关键的模块：

1.算子层维度推导：对于特定的算子，给定输入的维度信息，可以推导出输出的维度信息，即可在整个网络中完成维度的推导

2.维度约束化简：当完成维度的推导之后，网络中可能存在多个符号维度信息，在一个合理的网络中，部分算子是存在一些约束信息，比如 concat 算子，除 concat 的维度之外，要求其余的维度信息完全相等，通过挖掘这类算子的约束信息，能够对符号进行化简。

下面是一个简单网络进行维度推导和约束化简的示意图。

编译器后端主要负责将前端处理后的 IR 转换为目标硬件可执行的代码或硬件描述。主要功能包括基于硬件特性的 IR 优化、高效内存管理和代码生成等。下面将详细介绍各个模块。

ScheduleBlock(root){  serial for (i, 0, 32)  {    serial for (j_0, 0, 64)    {      serial for (j_1, 0, 128)      {        ScheduleBlock(A)        {          vi, vj = axis.bind(i, j_0 * 64 + j_1)                    A[vi, vj] = X[vi, vj] * 2        }      }    }  }}

CINN AST IR 中包含了以下信息，但集合和映射并不显示使用某种数据结构进行存储。

  • 访存关系：语句实例 <—> 内存单元

  • 依赖关系：语句实例 <—> 语句实例

  • 执行顺序：语句实例 —–> 语句实例

  • 执行顺序 = 语句实例的先后关系

语句实例集合范围 = 循环边界 + 循环步长 —— 循环构成一个带约束的整数空间，即迭代空间，迭代空间决定了语句实例，语句实例充满了迭代空间。

Schedule 为定义在 CINN AST IR 上的优化策略，常见的 Schedule 包括：LoopAlignment, Tile, Inline, Vectorize, Unroll 等。

以一个组合算子为例模拟可能的 AST 变换过程：

 [S1, S2, 1024] ==E=> [S1, S2, 1024] ==R=> [S1, S2] ==E=> [S1, S2] ==B=> [S1, S2, 1024] ==E=> [S1, S2, 1024]

serial for (i, 0, S1)  serial for (j, 0, S2)    serial for (k, 0, 1024)      ScheduleBlock(A)        vi, vj, vk = axis.bind(i, j, k)        A[vi, vj, vk] = X[vi, vj, vk] * 2serial for (i, 0, S1)  serial for (j, 0, S2)    serial for (k, 0, 1024)      ScheduleBlock(B)        vi, vj, vk = axis.bind(i, j, k)        B[vi, vj, vk] = A[vi, vj, vk] + 1serial for (i, 0, S1)  serial for (j, 0, S2)    ScheduleBlock(C__reduce_init)        vi, vj = axis.bind(i, j)        C_init[vi, vj] = 0serial for (i, 0, S1)  serial for (j, 0, S2)    serial for (k, 0, 1024)        ScheduleBlock(C)        vi, vj, vk = axis.bind(i, j, k)        C[vi, vj] = C[vi, vj] + B[vi, vj, vk]serial for (i, 0, S1)  serial for (j, 0, S2)    ScheduleBlock(D)      vi, vj = axis.bind(i, j)      D[vi, vj] = C[vi, vj] * 2serial for (i, 0, S1)  serial for (j, 0, S2)    serial for (k, 0, 1024)        ScheduleBlock(E)        vi, vj, vk = axis.bind(i, j, k)        E[vi, vj, vk] = D[vi, vj]serial for (i, 0, S1)  serial for (j, 0, S2)    serial for (k, 0, 1024)      ScheduleBlock(F)        vi, vj, vk = axis.bind(i, j, k)        F[vi, vj, vk] = E[vi, vj, vk] + 1

serial for (sp, 0, S1 * S2)    serial for (rb, 0, 1024)        ScheduleBlock(XXX)      vsp1, vsp2, vrb = axis.bind(sp / S2, sp % S2, rb)      XXX = XXXXXXserial for (sp, 0, S1 * S2)     ScheduleBlock(XXX)     vsp1, vsp2 = axis.bind(sp / S2, sp % S2)     XXX = XXXXXX

(3) Tile: 对所有 ScheduleBlock 的 loop nest 做相同的 Tile

// pure_spatial 轴 Tile 为：-1 * 16 * 64   Tile size 可为参数传入serial for (sp1, 0, S1 * S2 / 1024)  serial for (sp2, 0, 16)    serial for (sp3, 0, 64)     // S1 * S2 / 16 / 64, predicate: sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2      XXXXXX// impure_spatial_iter 轴 Tile 为 32serial for (sp1, 0, S1 * S2 / 1024)  serial for (sp2, 0, 16)    serial for (sp3, 0, 64)      serial for (rb1, 0, 32)        serial for (rb2, 0, 32)          ScheduleBlock(XXX)            predicate = sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2            vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)            vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)            vrb = axis.bind(rb1 * 32 + rb2)            XXX = XXXXX

// 例如 ScheduleBlock(A) inline 到 ScheduleBlock(B)serial for (sp1, 0, S1 * S2 / 1024)  serial for (sp2, 0, 16)    serial for (sp3, 0, 64)      serial for (rb1, 0, 32)        serial for (rb2, 0, 32)          ScheduleBlock(A)            predicate = sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2            vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)            vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)            vrb = axis.bind(rb1 * 32 + rb2)            B[vsp1, vsp2, vrb] = (X[vsp1, vsp2, vrb] * 2) + 1

(5) Reduce 优化: two step reduce&绑定部分 reduce 轴到 cuda

// 为了简洁，此处省略 reduce_init Block 和 predicateserial for (sp1, 0, S1 * S2 / 1024)  serial for (sp2, 0, 16)    serial for (sp3, 0, 64)      CudaBind[ThreadIdx.x] for (rb1, 0, 32)        serial for (rb2, 0, 32)          ScheduleBlock(C_rf)            vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)            vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)            vrb1 = axis.bind(rb1)            vrb2 = axis.bind(rb2)            C_rf[vsp1, vsp2, vrb1] = C_rf[vsp1, vsp2, vrb1] + B[vsp1, vsp2, vrb1 * 32 + vrb2]        ScheduleBlock(C)          vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)          vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)          vrb1 = axis.bind(rb1)          C[vsp1, vsp2] = C[vsp1, vsp2] + C_rf[vsp1, vsp2, vrb1]

(6) 循环融合: ComputeAt && SimpleComputeAt，融合外层循环乘积相同的循环，并且保证不破坏图级别依赖（规则负责）和元素级别依赖（原语负责）

serial for (sp1, 0, S1 * S2 / 1024)  serial for (sp2, 0, 16)    serial for (sp3, 0, 64)      ScheduleBlock(D)        vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)        vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)        D[vsp1, vsp2] = C[vsp1, vsp2] * 2      serial for (rb1, 0, 32)        serial for (rb2, 0, 32)          ScheduleBlock(E)            vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)            vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)            vrb = axis.bind(rb1 * 32 + rb2)            E[vsp1, vsp2, vrb] = D[vsp1, vsp2]          ScheduleBlock(F)            vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)            vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)            vrb = axis.bind(rb1 * 32 + rb2)            F[vsp1, vsp2, vrb] = E[vsp1, vsp2, vrb] + 1

(7) Bind Cuda 轴：在第二步中，所有 ScheduleBlock 对应的循环要 bind 到同一 Cuda 轴

serial for (sp1, 0, S1 * S2 / 1024)  CudaBind[BlockIdx.x] for (sp2, 0, 16)    CudaBind[ThreadIdx.y] for (sp3, 0, 64)      CudaBind[ThreadIdx.x] for (rb1, 0, 32)        serial for (rb2, 0, 32)          ScheduleBlock(XXX)

在业务场景中，模型结构具备多样化的特性，不同的shape和网络结构场景下，性能优化策略有差异，同时为了能够充分发挥硬件的特定，达到极致的性能表现，飞桨神经网络编译器引入自动调优模块。对于输入的 shape，通过对shape的自动分析，生成最优的 Schedule 变换策略，达到最优的极限性能。

Codegen 在 CINN IR AST 上做前序遍历，打印出对应硬件的指令，并通过硬件相对应的编译器（如 llvm、nvcc 等）进行编译得到可运行的函数指针，该指针会被封装到 `JitKernelOp“ 中用于后续执行器的解析执行。

a. 以函数定义为例子，cuda kernel func 和 x86 kernel func 的不同的是，cuda kernel func 会在函数名前增加__global__

针对 x86硬件，转义ir::_LoweredFunc_的代码如下：

void CodeGenC::Visit(const ir::_LoweredFunc_ *op) {  PrintFunctionDeclaration(op);   str_ += "\n";  ...  ...}

void CodeGenCUDA_Dev::Visit(const ir::_LoweredFunc_ *op) {  str_ += "__global__\n";         PrintFunctionDeclaration(op);   str_ += "\n";  ...  ...}

b. 在动态形状场景下，还会 codegen 出 infer shape function, infer shape function 的 CINN IR 会在 Bucket Lowering 中得到，转义过程复用的 x86 硬件的 codegen。infer shape kernel 如下：

function fn_exp_0_subtract_0_infer_shape (kernel_args, kernel_args_num, tensor_shape_args){  int64 S0 = cinn_get_value_in_cuda_kernel_args(kernel_args, 2)  {                infer_shape_set_value(0, 0, S0, tensor_shape_args)    infer_shape_set_value(0, 1, 256ll, tensor_shape_args)  }}

编译器生成的 Kernel 代码需要与深度学习框架执行器完成交互和集成才能最终运行起来，因此需要基于执行器的运行调度接口对编译器生成的 Kernel 进行封装。接入执行器后在运行时对于经过编译器处理的子图将执行 CINN 生成的 Kernel, 否则将执行常规的 PHI 算子 Kernel。

通过飞桨神经网络编译器的自动优化，我们在生成式推理模型上相比基础版本，性能提升30%；在科学计算场景（Nvidia modulus）上， 基于飞桨提供的高阶自动微分及编译优化技术，对比 PyTorch 模型训练性能整体领先60%。

同时飞桨会持续在模型和多硬件场景上探索神经网络编译器的能力，借助这种前沿技术给更多的业务带来价值。

7月至10月特设《飞桨框架3.0全面解析》直播课程，技术解析加代码实战，带大家掌握核心框架、分布式计算、产业级大模型套件及低代码工具、前沿科学计算技术案例等多个方面的框架技术及大模型训推优化经验，实打实地帮助大家用飞桨框架3.0在实际开发工作中提效创新！

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

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

如有变动，敬请谅解。

【算子融合策略源码】

https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc

【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#