Skip to content

hands on model tuning with TVM and profile it on a Mac M1, x86 CPU, and GTX-1080 GPU.

Notifications You must be signed in to change notification settings

JackonYang/hands-on-tvm

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

56 Commits
 
 
 
 
 
 
 
 

Repository files navigation

Hands on TVM

Step 0: TVM Overview

冯思远-Learning-based Deep Learning Compiler: b 站视频。

30 分钟,讲了为什么要做 TVM,& AutoTVM 的基本原理。

简明扼要,新手友好。质量很高。

冯思远是 TVM 的核心贡献者之一,上交大毕业。这个视频是给上交 2020 级的学生讲的。

Step 1: 陈天奇 - MLC 课程 Jupyter notebooks

对原 notebook 的改修

修改后的 code:陈天奇 - MLC 课程 jupyter notebooks

改动点:

  1. 基于 tvm 仓库的代码。原课程基于 mlc-ai-nightly 库,虽然稳定,但与最新的 tvm API 不一样。
  2. 精简注释 & 代码。
    • 目标:要点突出,适合复习、快速查阅代码。
    • 原课程,细节详细,更适合萌新学习 tvm。

备注:

  1. 课程使用 relax,还未合入 tvm 主分支。实测,切到 unity 分支后,可用。
  2. 代码走读发现,unity 分支,主要就是新增了 relax 相关的 5 万行代码。

推荐 3 个 notebook

2_tensor_program_abstraction.ipynb

  1. 用 TVM 的 2 种 DSL 写算子。
    • TVMScript 的 TIR 写 矩阵加法
    • Tensor Expression(简称 TE)写矩阵乘法
  2. 程序优化程序。tvm scheduler 把矩阵乘的性能提升 ~10+ 倍。

4_Build_End_to_End_Model.ipynb

  1. 用 TVM 手写模型。用到了 Relax 的 2 个 API: Dataflow, call_dps_packed / call_tir
  2. 可以注册外部算子,并在模型里混合使用。使用 call_dps_packed 调用外部算子。
  3. 把模型权重 bind 到 IRModule 上,运行时不用再传入模型权重。

5_Automatic_Program_Optimization.ipynb

  1. 这是 tvm 的精髓:autotvm & ansor。
  2. 使用 meta_schedule.tune_tir API,可以自动搜索最优的 schedule。
  3. 调用 tune_tir 时,可以手动指定 search space,但,如果写的太简单了,效果很可能不如默认版本。
  4. 可以把 Model 里的特定算子拿出来,单独调用 tune_tir 优化,并把优化后的算子,放回 Model。

备注:

  1. 原 notebook 花了大量篇幅讲采样搜索的最优化方法。实际,对应的 TVM API 很简单。熟悉最优化理论的,直接看 API 就够了。
  2. 搜索算法,依赖 xgboost。 pip install xgboost

其他章节

一共 10 节课。1 和 10 都是 ppt,没有 code。

其他章节的 code,个人看法如下

3_TensorIR_Tensor_Program_Abstraction_Case_Study_Action.ipynb

  1. 针对高性能计算新手的科普。原课程,花了大量篇幅讲解,为什么要 tiling, reorder。都是与 TVM 关系不大的背景知识。
  2. TVM API 的使用,基本是课程 2 的翻版。用的也是 TVMScript (TIR),简单讲了 TE。新东西不多。
  3. 特点是,例子是深度学习模型里的基础结构:线性层 + ReLU 激活。

整体来讲,信息量不大,跳过。暂不细致整理。

6_Integration_with_Machine_Learning_Frameworks.ipynb

  1. 个人理解,属于 TVM 定制开发,而且是针对 relax 前端的。普通用户不需要。
  2. 现在用的都是 relay,前端的算子注册跟 relax 也不一样。学会了也没啥用。
  3. 如果是新手,通读一下,了解图前端如何转模型和算子的,还不错。
  4. 注意:高性能算子的开发和注册,跟这个不一样。要写 Tensor Expression 和 TVMScript,见 2、3 节。

7_GPU_and_Specialized_Hardware.ipynb

  1. 值得看。讲了 TVM 如何支持 GPU。
  2. 花了大量篇幅介绍 GPU。实际 TVM API 相关的,内容应该很少。
  3. 手头的 GPU 环境没配,先跳过。

8_GPU_and_Specialized_Hardware_part2.ipynb

  1. 讲的是 NPU 和 FPGA。
  2. 个人的经验,坑还是很多的,而且有些坑是 TVM 架构难以处理的。完全不像陈天奇说的这么简单。
  3. 华为海思使用 TVM,据说是一二百人做了两三年才开始产生商业价值。

9_Computational_Graph_Optimization.ipynb

  1. 讲的主要是如何用 pattern match and rewrite 做算子融合。
  2. 此处的算子融合,还是基于规则的。信息量不大。
  3. pattern match and rewrite 是编译器的基本功。没做过编译器的话,可以练练手。

Step 2: TVM 自动优化模型 - 官方 docs 精简版

说明:

  1. 代码均来自 TVM 官方文档的 User TutorialHow To Guides
  2. 稍作修改和整理,目标:要点突出,适合复习、快速查阅代码。
  3. 更详细的使用方法,需要看官方的文档、源码。
  4. 出于学习的目的,使用了比官方文档更小的模型 & 搜索空间。demo 不代表 AutoTVM 和 TVM Ansor 的优化效果差异。
  5. Mac M1 上,AutoTVM 和 Ansor 都抛异常,与 rpc 有关,具体原因未分析。个人怀疑,如果熟悉 tvm 和 rpc 协议,应该是个不太难修的问题。

01-build-run-dl-models-tvm.ipynb

茴香豆的茴有几种写法? -- 跟 TVM 编译模型的方法一样多。

  1. TVM build & run 模型,有多种 API 组合可选。粗读代码发现,一套 API 组合对应一套轮子。
  2. 功能基本一样,用起来也基本一样。我感觉,掌握一套即可。
  3. 我选了 relay.build + tvm.contrib.graph_executor,可以 export_library API 把编译后的模型保存为 .so 文件,方便 deploy 分发。

02-tvm-ansor-template-free-tune.ipynb

  1. 使用 ansor (tvm.auto_scheduler) 自动优化 AlexNet 模型。
  2. 搜索 5min,模型加速 2x。硬件: Intel(R) Xeon(R) Gold 5320 CPU。
  3. 新人友好,上手简单。
    • ansor 是 template free 的优化器,不要求理解模型和硬件。
    • 模型优化相关的代码,一共不到 10 行代码。

03-autotvm-tune-with-templates.ipynb

  1. 使用 autotvm 自动优化 AlexNet 模型。
  2. 搜索 13min,模型加速 5x。硬件: Intel(R) Xeon(R) Gold 5320 CPU。
  3. 原理上需要写 template,但 CPU/GPU 上,使用默认 template,无脑暴力搜,结果也不错。
  4. 代码量比 ansor 稍微多一点,但也非常简单。

TODO:

  1. dynamic shape, relax, unity & Nimble
  2. https://github.com/mlc-ai/mlc-llm
  3. TVM Papers: https://tvm.apache.org/docs/reference/publications.html

optional:

  1. microTVM: TVM on bare-metal
  2. VTA: Versatile Tensor Accelerator

About

hands on model tuning with TVM and profile it on a Mac M1, x86 CPU, and GTX-1080 GPU.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 99.9%
  • Python 0.1%