Skip to content

doongz/mlc-ai

Repository files navigation

Machine Learning Compilation

Introduction

  • 所属大学:个人公开课
  • 授课老师:陈天奇
  • 先修要求:一定的深度学习框架背景知识,有系统层面的编程经验
  • 编程语言:Python
  • 课程难度:🌟🌟🌟
  • 预计学时:30 hours
  • 学年:Summer 2022

这门课是机器学习编译领域的顶尖学者陈天奇在2022年暑期开设的一门在线课程。其实机器学习编译无论在工业界还是学术界仍然是一个非常前沿且快速更迭的领域,国内外此前还没有为这个方向专门开设的相关课程。因此如果对机器学习编译感兴趣想有个全貌性的感知的话,可以学习一下这门课。

本课程主要以 Apache TVM 这一主流的机器学习编译框架为例(陈天奇是这个框架的创始人之一),聚焦于如何将开发模式下(如 Tensorflow, Pytorch, Jax)的各类机器学习模型,通过一套普适的抽象和优化算法,变换为拥有更高性能并且适配各类底层硬件的「部署模式」。课程讲授的知识点都是相对 High-Level 的宏观概念,同时每节课都会有一个配套的 Jupyter Notebook 来通过具体的代码讲解知识点,因此如果从事 TVM 相关的编程开发的话,这门课有丰富且规范的代码示例以供参考。

所有的课程资源全部开源并且有中文和英文两个版本,B站和油管分别有中文和英文的课程录影。


端侧 ai 的一门课,非常契合业界使模型落地的业务

mlc_process

机器学习编译的「四大抽象」:计算图、张量程序、算子库、硬件指令

机器学习编译的「核心思想」是:给定一系列的张量函数,对它做一系列的变换,变换成一些等价的但更优化的,在对应平台上更有利于实现,这些变换是由程序自动完成。

机器学习编译本身的「开发模式」和传统的开发模式有些不同的地方是,会强调变换 Transform ,可以通过 tvm 编写、tensor expression构造、代码自动生成拿到一个初始的函数,它是的数据结构是 IRModule 来表示,也就是开发到一个共同的抽象 IRModule。接下来,往往对 IRModule 的张量函数做一系列的变换和优化。几乎所有的机器学习的编译都是对 IRModule 的变换,变换之后拿到另外一个版本,然后调 build 函数,变成可部署的形式

这门课程的核心:

  • 讨论用哪些方式表示不同的张量函数在IRModule里面;
  • 会有哪些变换,可以使得原来的函数变换成在一些环境下更优的函数;

Resources

Notes

  • 机器学习编译的目标
    • 集成与最小化依赖
    • 利用硬件加速
    • 通用优化
  • 为什么学习机器学习编译
    • 构建机器学习部署解决方案
    • 深入了解现有机器学习框架
    • 为新兴硬件建立软件栈
  • 机器学习编译的关键要素
    • 张量和张量函数
    • 抽象和实现是值得思考的工具
  • 元张量函数(单算子)表示机器学习模型计算中的单个单元计算。

    • 一个机器学习编译过程可以有选择地转换元张量函数的实现。
  • 张量程序是一个表示元张量函数的有效抽象。

    • 关键成分包括: 多维数组,循环嵌套,计算语句。
    • 程序变换可以被用于加速张量程序的执行,并行计算。
    • 张量程序中额外的结构能够为程序变换提供更多的信息。
  • TensorIR 抽象

    • 包含循环、多维缓冲区等常用元素
    • 引入了一个封装循环计算要求的新结构
    • 可以在 Python AST 中构建(通过 TVMScript)
  • 我们可以使用变换来创建不同的 TensorIR 变体,通过并行化、向量化与循环展开等。

  • 通用 MLC 流程:开发、变换、构建。

  • 计算图抽象有助于将元张量函数拼接在一起以进行端到端执行。
  • Relax 抽象的关键要素包括
    • call_tir 构造,将目标传递规范的元函数嵌入到计算图中
    • Dataflow block
  • 计算图允许调用环境库函数和 TensorIR 函数。
  • 将参数绑定到 IRModule(模型中存放权重)
  • Build
  • 「随机变换」帮助我们指定可能程序的搜索空间。
  • Meta-Schedule 在搜索空间中搜索,并找到优化后的程序。
  • 我们可以使用另一种变换,将初始的元张量函数替换为优化后的函数,并更新端到端执行流程。

算子搜索:“可能的变换是什么”。

  • 更多级的循环转换
  • 中间计算的矢量化
  • 并行化和循环展开

分析程序相对应的一些规律,尝试一些变换

如何进行搜索(机器学习最优化算法),如何为搜索空间建立规则(block大小,算子间的关联性等),可以借鉴CS50’s ai 这门课

  • 张量表达式 API 允许我们创建原始的 TensorIR 函数。
  • BlockBuilder API 通过 emit_te 和其他函数创建 IRModule。
  • 通过将模型转换为 IRModule,实现与现有的机器学习框架的整合。
  • 典型的 GPU 包含两级层次结构。 每个线程由(在 CUDA 术语中)threadIdx.xblockIdx.x 索引(也可以有多个维度索引,但它们可以融合为一个)。

  • 共享内存有助于缓存同一块内的线程中常用的数据。

  • 在 GPU 优化期间鼓励内存重用。

  • 硬件专业化向张量计算的总体趋势。

  • 带有张量化 block 的 TensorIR 变换。

  • 张量化:将循环计算 block 映射到专门实现的过程。

  • 我们可以通过改写计算图数据结构来优化模型。
  • 使用访问者模式改写调用节点。
  • 我们可以进行计算图转换,例如融合和循环级代码生成。

Releases

No releases published

Packages

No packages published