Intership｜Paddle的amp pass（混合精度 pass）新旧迭代中遇到的问题

缘起

一直都想要参加一些大项目的开源活动，同时自己也想在 AI inference 上获得更深入的了解，一直在寻找这样一个可以参与的机会。机缘巧合之下在今年参与了 paddle pass 组件的开发，虽然只是其中一个很小的 pass，但是在其中也是收获满满～

paddle现在的 amp pass还不支持 PIR 的直接使用，就需要经过下面的流程：

image-20241028225011825

对于一个新生的 PIR 来说，一直使用旧的 pass 去做转化显然是不理想也不能接受的，于是从支持旧 IR 的 pass 到支持PIR 的 pass 就要有一个转化的工作。旧 IR 从一种相对落后的实现迭代到现在的 PIR，有许多工作需要补充，就像 paddle 说的，“PIR灵活可扩展，在底层抽象了一套高度可扩展的基础组件，涵盖Type、Attribute、Op、Trait和Interface，并引入了Dialect的概念，赋予开发者灵活扩展与自由定制的能力，从而提供了全面且稳健的语义表达能力。在模型表示层，通过多Dialect的模块化管理和统一多端表示，实现了训练与推理一体化的全架构统一表示，实现了算子和编译器的无缝衔接，支持自动优化和多硬件适配。在图变换层，通过统一底层模块并简化基础概念，向用户提供了低成本、易用且高性能的开发体验，以及丰富且可插拔的Pass优化机制。飞桨PIR坚守静态单赋值（SSA）原则，确保模型等价于一个有向无环图，并采用Value和Operation对计算图进行抽象，其中Operation代表节点，Value代表边”。

在我实际使用中，也的确发现了 PIR 相较于旧 IR 的很多优势。如果 amp pass 可以直接工作在PIR 上，那么流程就会变为下面这样：

image-20241028230357527

对于这套流程来说，最终态就是从 PIR model 直接经过 PIR pass 送入执行器推理。

过程

旧 IR 的 amp pass 执行流程如下：

image-20241028230622643

在使用了PIR 后，其中的一些步骤就可以省略。由于PIR 的 SSA 特性，每 个变量只有唯一的赋值点， 所以我们可以把用来设置 Op 唯一标志的这步省略。由于PIR 中引入的这一套Value 和 Operation 的抽象，我们在设置权重和输入输出的精度时不用 在分开做，插入 cast op 也就不 需要再等前两个做完再做。融合 到 ProcessBlock 里，通过修改 op 的attribute 里的精度相关设置以及，设置输出中的 Type 来完成相关功能。PIR amp pass 就变成了下面这样（灰色是可以融合或者删除的步骤）：

image-20241028230809409

对于 PIR 的 pass 来说，不需要关心上层的实现机制，只需要利用 paddle pir pass 的两种开发机制：DRR 和 Pattern rewrite

image-20241028231427016

PR 机制只需要可以选择直接继承 pir::Pass，然后自定义接口 Run 去实现匹配功能，比如下面官方给出的常量折叠 pass：

image-20241028231642830

由此可见，amp pass 的开发也可以借鉴这一点，只需要聚焦 pass 的顶层实现，由于目前的 ir 可能会存在一种循环嵌套的可能性，由此 pass 的实现可以如下,采用Pattern Rewriter机制 ，只要实现CanApplyOn 就可以 End 控制 Run 函数访问到最外层的 op，然后进入后通过递归处理就可以就可以实现对所有 op 的深度优先遍历，解决多层循环嵌套的问题。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
void Run(pir::Operation* op) override {
    SubBlockRun(op->GetParentProgram()->block());
    AddStatistics(op_run_low_precision_.size());
  }

bool CanApplyOn(pir::Operation* op) const override {
  return op->num_regions() > 0 && op->isa<pir::ModuleOp>() && place_ == paddle::PlaceType::kGPU &&
         (precision_mode_ == phi::DataType::FLOAT16 ||
          precision_mode_ == phi::DataType::BFLOAT16);
}
  
void SubOpRun(pir::Operation* op){
  for (auto& region : *op) {
    for (auto& block : region) {
      SubBlockRun(&block); //subblock
    }
  }
}

void SubBlockRun(pir::Block* block){
  GetOpPrecision(block);
  UpdateOpPrecision(block);
  pir::Builder builder = pir::Builder(context_, block);
  ProcessBlock(block, builder);
  cached_cast_ops_.clear();
  op_should_not_handle_.clear();
  for (auto& op : *block) {
    SubOpRun(&op);
  }
}

在完成开发工作之后，还参与到了该 pass 的性能测试中。主要是利用下面这种形式展开性能测试，用paddle自己的 profile 工具，生成可转化的 json 文件，去导出 timeline 、转化成 mysql 去分析。最初没有发现 perfetto 的转化db 文件的功能， 但是后面了解到，用 sql 工具处理数据可以得到很神奇的效果。后面也想补充一篇关于 perfetto 使用的技巧，互联网上关于这个真的很缺，ai inferencer 真的很伤。

image-20241028232128061

通过这种方式，发现了在使用 paddle-inference-demo 的使用问题并解决，最后得到了下面比较好的 tipc测试结果。

image-20241028234113104

最终在多次代码更新加性能测试后，初步“收敛”了这个 amp pass。

去往

这次的工作对我自己的研究生路途也有很重要的启发：

  • 一种优秀的 IR 对于训推工作来说有巨大的便利性，便利性就可能体现在代码编写上或者调试中。在以往我常用的推理框架中不具备这种特性导致推理流程进行的不是很顺畅。

  • 接触到 paddle推理的开发，让我对整个推理流程有了新的认识，推理需要精心设计数据的流动，最终在更优的计算资源和时延控制下，产生精确的结果，让算法真正从理论走向实践。

  • 优秀的开源项目的确存在一种传承关系，远乐导师的帮助让我快速的融入这个课题，更是让我对 paddle 推理相关有了更深的认识，未来会有更多机会加入到优秀的项目之中。

Code: https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/pir/transforms/general/auto_mixed_precision_pass.cc