计算图优化

模型导出后做推断目的

程序编译执行过程

计算机就是一层一层的抽象, 一层包一层. 越上层越简洁, 但是性能越差. 从上到下的全局理解很难, 但是对于实际性能优化很重要.

DEMO HERE

Python程序优化套路: Python写意图, C改写实现

Python for clarity, C for performance

最简单的CPU后端:

编译为C的方式

计算图

模型 = 结构 + 参数

DAG

算法模型=计算图

不涉及IO, 不产生副作用, 一定固定步数执行完成输出的, 纯函数

静态图 VS 动态图

静态图

动态图

主观感受: 是否支持动态尺寸输入

殊途同归:

torch.jit / TorchScript IR

PyTorch模型代码的IR形式, 运行时可不依赖Python环境, 但是依赖一个较大的SDK (LibTorch)

直接用pickle做序列化格式

基本上Python语法的子集: https://pytorch.org/docs/stable/jit_language_reference.html

torch.jit.script

静态代码解析提取, PY代码翻译为计算图, 保留了动态控制分支流

torch.jit.tracing

需要提供样本输入, 反射机制跟踪计算过程, 只看到实际执行的路线, 参数常量化, 抹掉动态控制流, 执行性能更好

@torch.jit.script 手动声明需要导出/忽略或者包装的方法

DEMO

PyTorch 2.0

TorchScript的问题: 需要代码编写者做特别处理, 额外的心智负担

主要卖点是提速

torch.compile 默认情况尽可能把可以不用Python执行的环节抠出来编译执行, fullgraph全部导出为计算图, 但是对于代码写法有严格约束

怕ONNX抢饭碗???

torch.compile主观对比感受 (BERT分类任务)

4090 上一定提升 2080 上负优化

GPU优化针对的”高端”服务器显卡, “低端”RTX游戏卡似乎和感知不大

目前局限: 不支持PY3.11 / 不支持Window

ONNX / Open Neural Network Exchange

ONNX: 框架无关的IR, 或者说数据标准/协议

微软撑腰的项目

ONNX Runtime (ORT): 跨平台的ONNX模型推断(及训练)加速器

计算图捕获: torch.onnx.export利用torch.jit.trace后导出ONNX文件

导出的计算图是未处理的 (除了常量折叠), 推断时根据执行环境对计算图优化

有针对模型结构的优化器

DEMO TIME

PyTorch模型 153QPS / BERT优化后ONNX 360QPS

动态尺寸 / dynamic axis

batch / 批量大小 / 一般推断场景单批就行了, 不必要动态batch

序列模型场景, 需要申明动态尺寸

动态尺寸会导致不能执行一些优化, 但是从性能上来说比补零定宽更值得

计算图优化

常量折叠 / Constant Folding

算子融合 / Fusion

加乘算子 / Matmul Add Fusion

BERT Embedding Layer Fusion

…

数据布局重排 / Layout optimization

NCHW for CNN

卷积计算转换为两矩阵相乘

并行计算

NOTE 并行 (parallelism) vs 并发 (concurrency) 区别

硬件指令执行层面并行

提升速度手段:

e.g. AVX

单算子并行

CPU上是OpenMP编程, 映射到指令层面并行

def mul(a, b, n):
    #pragma omp parallel for
    for (i = 0; i < n; i++) {
        a[i] *= b[i]
    }

多算子并行 / 子图分解

计算图没有依存关系的部分可以同时计算

类比: 运筹学 / 调度任务安排

线上推断运行问题: OpenMP错误配置导致实际推断性能下降

容器化推断环境运行, 默认拿到的是宿主机器核数, 实际上限制了CPU数, 错误的安排线程数, 反而会导致性能下降

需手动基于容器配置感知修改OMP_NUM_THREADS, 或者改session_option

计算图训练优化

计算图加速不光对于推断(前向过程)有意义, 对于训练(+后向过程)同样意义重大

分布式训练 / 混合精度训练 / …

推断/训练需求算力越来越大的今天, 关注计算图优化的意义: 小我: 努力省钱, 将本增效; 大我: 低碳排放, 保护地球 !

实践中训练最有效的优化措施: 提高数据读写IO (磁盘到内存, 内存到显存), 把GPU利用率打满, 把显卡烤熟后再考虑其他手段

努力编写可导出的模型代码

现状: 训练时torch逻辑, 推断时numpy逻辑, 写双份.

期望: 预处理/后处理尽可能打包到模型中, 避免逻辑同步维护及正确性校验工作, 确保训练推断一致; 以及能显著优化推断性能

要相信, 只要不涉及外部IO的, 都可以翻译为计算图逻辑; 导出的再挫, 也一定比手搓的PY代码要强

例子

https://onnxruntime.ai/docs/reference/operators/add-custom-op.html

程序优化指南

其他模型推断加速手段

Reference

HOME