愛鋒貝

標(biāo)題: 模型加速與 AI compiler 介紹 [打印本頁]

作者: 毛屌屌 時間: 2023-4-4 09:03
標(biāo)題: 模型加速與 AI compiler 介紹
目標(biāo)

非從業(yè)者：了解 ai compiler 的研究內(nèi)容、成本和意義。
AI 從業(yè)者：幫大家讀了一些經(jīng)典的 paper。
主要討論 deep learning 模型的推理（inference）。

AI 從業(yè)者的定義：深度參與過模型的訓(xùn)練優(yōu)化或推理部署。
不涉及什么

系統(tǒng)的介紹一門學(xué)科、理論。
聽完就能干活、能學(xué)以致用。

背景介紹

當(dāng)我們談到深度學(xué)習(xí)時，可能在談什么

(, 下載次數(shù): 172)
出自: https://www.modb.pro/doc/48500 飛槳：源于產(chǎn)業(yè)實踐的開源深度學(xué)習(xí)平臺 -藍翔 _百度
模型訓(xùn)練的速度 - from paper

Model Name	Training Time	時間	Hardware	Data size	Paper Link
Transformer	12h	2017.06	8 P100 GPU	37000 token	http://arxiv.org/abs/1706.03762
BERT	81.4h	2018.10	16 TPU	3.3B word corpus	https://arxiv.org/abs/1810.04805
BERT	76 min	2019.4	1024 TPU	3.3B word corpus	https://arxiv.org/abs/1904.00962
XLNet	2.5days	2019.6	512 TPU v3 chips	32.89B	https://arxiv.org/abs/1906.08237
Resnet50	2.2 min	2018.11	TPU v3 Pod	ImageNet	https://arxiv.org/abs/1811.06992
Resnet50	75s	2019.3	2048 GPU v100	ImageNet	https://arxiv.org/abs/1903.12650
GPT	month	2018.6	8 GPU	BooksCorpus 800M words	https://www.cs.ubc.ca/~amuham01/LING530/papers/radford2018improving.pdf
GPT-2	week	2019.2	256 of Google&#39;s Cloud TPU v3	23 million URLsover 10 million HTML pages	https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf
BigGAN	24~48h	2018.9	TPU v3 Pod	ImageNet	https://arxiv.org/abs/1809.11096
ResNet 10132*16D	22days	2018.5	336GPU	3.5 billion images	https://research.fb.com/wp-content/uploads/2018/05/exploring_the_limits_of_weakly_supervised_pretraining.pdf
RoBERTa	1day	2019.7	1024 v100	4倍 XLNet, 40倍 BERT	https://arxiv.org/pdf/1907.11692.pdf
ELMo	2weeks	2018.2	3 GTX 1080	5.5B tokens	https://arxiv.org/abs/1802.05365

隨著預(yù)訓(xùn)練技術(shù)的普及，從頭訓(xùn)練模型的場景越來越少了。
以 bert 為例，機器翻 64 倍，訓(xùn)練時間下降 64 倍。-- 砸錢就能加速。這是一個分布式的問題。

模型訓(xùn)練 vs 推理部署

深度學(xué)習(xí)框架的數(shù)量，那可以太多了。
訓(xùn)練框架

TensorFlow, Google
Pytorch, FaceBook
MxNet, Amazon
Caffe, Caffe2
CoreML，蘋果，代碼保密
Paddle, 百度。徐偉，百度第一位 T11。后任地平線 AI 首席科學(xué)家。
MegEngine, 曠視
Mindspore, 華為
OneFlow, 一流科技

推理框架

OpenVINO, Intel
ARM NN, ARM
TensorRT, nvidia
ONNXRuntime. ONNX 節(jié)點粒度較細(xì)，推理速度有時候比其他推理框架慢.
ncnn(騰訊) 手機端推理框架。
MNN（阿里）手機端推理框架。
MMDeploy。 OpenMMLab 的框架。

國產(chǎn)框架：

Jittor: 清華。PR 稿里見得多。
Paddle: 百度
OneFlow: oneflow 公司。

Mindspore: 華為

OpenMMLab: 香港中文大學(xué)-商湯科技聯(lián)合實驗室

深度學(xué)習(xí)框架簡史：https://syncedreview.com/2020/12/14/a-brief-history-of-deep-learning-frameworks/
總結(jié)要點：

AI != 模型訓(xùn)練。推理市場很大。
框架的核心需求：易用性高、性能好。
推理和訓(xùn)練，基本是 2 套框架了。這就是“專業(yè)”。
訓(xùn)練強調(diào)開發(fā)效率（易用性），比如 tensorflow，pytorch。性能優(yōu)化的技術(shù)棧偏分布式的那一套。
推理強調(diào)執(zhí)行效率（性能）。比如 TensorRT，OpenVINO。關(guān)注性價比，最終歸宿都是與硬件高度綁定。

模型部署的軟件棧

百度 PaddlePaddle 的工具鏈
出自: https://www.modb.pro/doc/48500

(, 下載次數(shù): 165)

(, 下載次數(shù): 184)
Paddle 賺錢

(, 下載次數(shù): 141)
華為 MindSpore
來自 https://www.slideshare.net/Huawei/introduction-to-huaweis-fullstack-ai-portfolio

(, 下載次數(shù): 156)

(, 下載次數(shù): 169)

(, 下載次數(shù): 153)
來自 https://www.mindspore.cn/tutorial/zh-CN/r1.2/introduction.html

(, 下載次數(shù): 164)
寒武紀(jì)軟件棧介紹

(, 下載次數(shù): 166)
Cambricon PyTorch 訓(xùn)推一體：系統(tǒng)架構(gòu)分為前端和后端兩部分，前端指的是 Torch 和 Catch 中的各種算子接口，后端指的是 CNNL 和 MagicMind ；
CNNL：寒武紀(jì)人工智能計算庫，支持豐富基本算子、可變輸入推理、量化訓(xùn)練和混合精度訓(xùn)練；
MagicMind：寒武紀(jì)推理加速引擎，支持靈活輸入維度、多種量化模式、多種圖優(yōu)化算法，給用戶帶來極致推理性能；
高性能計算入門

軟件工程師需要了解的硬件特性
要點總結(jié)：

計算單元，并行能力越來越強。標(biāo)量 -> 向量 -> 張量。見下圖。
存儲，越快越貴，多級 cache。

From https://zhuanlan.zhihu.com/p/613611390
從 CPU 到 GPU 是 2015、2016 年左右，從 GPU 到 TPU 是 2019 年左右。

(, 下載次數(shù): 191)
優(yōu)質(zhì)讀物

高性能計算的研究內(nèi)容
高性能計算，主要針對科學(xué)計算。研究的核心是如何加速 nested loops。
科學(xué)計算的特點，是大量的數(shù)據(jù)運算 -> 大量 nested loops
優(yōu)化 nested loops 的手段，主要是并行計算。
并行計算分 4(or many) 類:

位級（bit）-- CPU 指令經(jīng)歷了 8bit, 16bit, 32bit, 64bit。
ILP 指令級 (instruction level Parallel)
以 RISC 為例，先做了流水線，然后是多指令發(fā)射。
傳統(tǒng)編譯器和硬件都是隱式的 ILP。
DLP 數(shù)據(jù)級 (data level Parallel)
比如：SIMD (單指令多數(shù)據(jù)), GPU。常見于圖像處理、深度學(xué)習(xí)。
AI compiler 在這一層也發(fā)力。
任務(wù)級 (task) ?？紤]多處理器，內(nèi)存模型很關(guān)鍵。共享內(nèi)存，或分布式內(nèi)存。分 2 個子類
TLP 線程級 (Thread level Parallel)
RLP request 級

前 3 類的優(yōu)化，非?；A(chǔ)，需要懂 hardware & Linux kernel。
第 4 類，比較常見。其中，分布式內(nèi)存的，通俗的說，就是分布式計算了。
高性能計算入門的 Hello world：矩陣乘 (General matrix multiply)，簡稱 GEMM。
加速的方法論
維基百科 - Amdahl&#39;s law
在問題總量不變的前提下，提升并行度帶來的加速收益，越來越小。

(, 下載次數(shù): 173)
維基百科 - Gustafson&#39;s law
Amdahl&#39;s Law 假定了問題的規(guī)模（計算量）不變，實際上，隨著資源利用率的改善，工程師們總是會增加新的計算需求進來，提升了問題規(guī)模。
因此，不能被并行優(yōu)化的部分，隨著問題規(guī)模的增長，對系統(tǒng)總體速度的影響，越來越小。
Roofline Model paper
一套算法，用于評估：
計算量為 A 且訪存量為 B 的模型在算力為 c 且?guī)挒?D 的計算平臺所能達到的理論性能上限E是多少。

(, 下載次數(shù): 165)
streetlight 方法
A policeman sees a drunk looking under a streetlight,
and asks what he is looking for.
The drunk says he has lost his keys.
The policeman can&#39;t find them either,
and asks if he lost them under the streetlight.
The drunk replies:
&#34;No, but this is where the light is best.&#34;
總結(jié)：只做會的，不做對的。
Why AI compiler

模型加速。適配特定硬件的加速，省出來的算力都是錢。
易用性、硬件 bug 遮羞布。AI 芯片廠做 ai compiler 基本是標(biāo)配。
除了 Nvidia，如果不是開箱即用，基本拿不到用戶。
早期的幾代芯片，bug 比 feature 不知高到哪里去了。ai compiler 就是最好的遮羞布。
摩爾定律終結(jié)，DSA 會越來越多，軟件棧人力需求將爆炸。趨勢：融合現(xiàn)有的編譯器設(shè)計快速的打造軟件生態(tài)。https://arxiv.org/abs/2002.11054 MLIR: A Compiler Infrastructure for the End of Moore&#39;s Law
加速之外，ai compiler 能講故事的場景也不少。
on-device learning。比如，iPhone 上 fine tune 自己的人臉開鎖模型，滿足數(shù)據(jù)安全要求。
DL 框架太多，學(xué)不動了。用工具鏈自動轉(zhuǎn)成一個框架的格式，學(xué)一個就夠了。

問題：AI compiler 會不會成為下一個被 AI 革了命的領(lǐng)域？
深度學(xué)習(xí)模型的加速手段

模型的表示 - DAG

Deep learning 模型就是一個 DAG 圖 Directed acyclic graph，有向無環(huán)圖。
建模方法：
Node: operator (運算，粗略理解，就是函數(shù)）
Edge: Tensor (張量，即，數(shù)據(jù)。粗略理解，就是函數(shù)的輸入?yún)?shù)和輸出數(shù)據(jù)，一個函數(shù)的輸出，同時又是下一個函數(shù)的輸入）
可視化工具：https://github.com/lutzroeder/netron
另一個古老的 caffe 工具：https://dgschwend.github.io/netscope/#/preset/inceptionv4

(, 下載次數(shù): 163)
如何讓模型跑的更快

工程上，就是 2 個 layer 的工作：

圖層。比如，
多算子融合后，提高了局部性。中間結(jié)果無需反復(fù)搬入搬出。
找 better 拓?fù)湫颍岣卟⑿卸取?br />
算子層。
單個算子跑得快。比如 GEMM 類的矩陣乘加速，算子的指令重排等。
用傳統(tǒng)編譯器的那套性能優(yōu)化方法。

編譯器優(yōu)化 loop optimization 三板斧:

Fusion
Tiling
vectorization

算法上，深度學(xué)習(xí)圈研究的比較多：

模型量化
稀疏化
模型蒸餾等。

原理：
模型中只有小部分參數(shù)起了核心作用，其他的大部分參數(shù)是無效參數(shù)，可以去除掉。即，矩陣稀疏化。
不能去掉的參數(shù)中，很多參數(shù)在推理時的狀態(tài)空間非常少，用 int8 就足以區(qū)分開。不需要 fp32。即，量化。

(, 下載次數(shù): 153)
Transformer 類模型的加速

加速方法，詳見 https://scale.com/blog/pytorch-improvements
2 個自測題：

從 LSTM 到 transformer 的例子。為什么 transformer 能瞬間霸榜 NLP，但霸榜 CV 就慢了。
LLM（NLP 大模型）輸出的 token （文字）數(shù)量與算力成本的關(guān)系。

Transformer 架構(gòu)的算力特征
Attention Is All You Need Arxiv
【李沐精讀】https://www.bilibili.com/video/BV1pu411o7BE
比論文更好讀的 blog: The Annotated Transformer
Transformer 架構(gòu)圖：

(, 下載次數(shù): 157)
作為對比，LSTM 架構(gòu)圖

(, 下載次數(shù): 162)
CNN 架構(gòu)

(, 下載次數(shù): 152)
CNN 和 transformer，都用一組結(jié)構(gòu)相同但 weight 不同的 block 去學(xué)習(xí)，
block 之間相互不依賴，可以 parallel 執(zhí)行。
LSTM 為代表的 RNN，都是 sequential 的處理，每一個 block 都要接收前序 block 的 output，存在數(shù)據(jù)依賴，很難并行。
詳細(xì)解釋：https://voidful.medium.com/why-transformer-faster-then-lstm-on-generation-c3f30977d747
總結(jié)：
LSTM 一直被算力鎖死了學(xué)習(xí)能力，換成 transformer 以后，學(xué)習(xí)能力就碾壓了。
CNN 被算力鎖死的比較輕，原本就發(fā)展的很厲害了。
用 cache 加速 transformer
實現(xiàn)參考:
https://github.com/alex-matton/causal-transformer-decoder
https://github.com/hpcaitech/CachedEmbedding
名詞解釋：decoder 時，每生成 1 個輸出 token 的過程，叫 timestep。

encoder output 和 decoder 可以并行計算。
Decoder 時是自回歸的，如果能做好 cache，那么，每個 decode timestep 只需要
計算新 token（上一個 timestep 的輸出）的 embedding。
計算新 token 和其他 token 的 attention

Input sequence length M, decode N tokens。
時間復(fù)雜度，input * output self-attention + output * output self-attention
優(yōu)化前: O(MN2 + N3)
優(yōu)化后: O(MN + N2)
實驗結(jié)果
We compare our three different implementations

The most naive Pytorch implementation (defined in the first piece of code), which uses nn.Transformer
The Pytorch encoder-decoder implementation (second piece of code).
Our CausalTransformerDecoder (third piece of code).

(, 下載次數(shù): 177)
深入 AI compiler - TVM

詳細(xì)介紹：https://mlc.ai/summer22/schedule
TVM Stack

AI compiler 如何理解模型？
TVM 的方法是：分 4 層 IR

(, 下載次數(shù): 172)
From https://zhuanlan.zhihu.com/p/613611390
TVM 的圖優(yōu)化

(, 下載次數(shù): 173)
解釋：

算完 3 就可以立即釋放 2
算 3 之前就算 5，則 2 和 5 就會同時 alive 在 cache 里
算完 3 再算 2，就不存在 2 和 5 同時 alive

整體思路：

本質(zhì)是找最優(yōu)的 DAG 的拓?fù)湫颉?br />
因為 cache 分級，如何把更多的計算放在更快的 cache 上，是主要目的。
因此，最佳拓?fù)湫虻哪繕?biāo)函數(shù)：最小峰值內(nèi)存占用

以 TFLite 用的方案為例，
參考 Paper: https://arxiv.org/abs/2001.03288 2020-google-efficient-memory-management.pdf
研究的問題：Mobile and embedded device: inferior physical memory
目標(biāo)：Smaller memory footprint
結(jié)果：11% smaller memory footprint than the state-of-the-art
弊端：

對動態(tài) shape 的處理，可能是不太行。
沒有考慮 cache hit rate 的影響。

建模的核心抽象，完整的 pipeline 分為如下 4 個 stage，每個時刻 t 可以執(zhí)行 1 個 stage。
- L(load) : load data from LLB->L1
- R(Compute): compute on L1
- S(Store): store result from L1->LLB
- F(Free): free data on L1
優(yōu)化理論的 2 大流派：

基于規(guī)劃理論：（整數(shù)）線性規(guī)劃，動態(tài)規(guī)劃、二次規(guī)劃和凸優(yōu)化等。
基于搜索算法：基于采樣理論、啟發(fā)式算法等。全局最優(yōu)、非全局最優(yōu)。
（流派 3）啟發(fā)式的策略，即，人工手動調(diào)度。

自動駕駛的 motion planning，路線之爭的時間線：

搜索為主。如果走規(guī)劃路線，是一個高維的非線性規(guī)劃問題，算力搞不定。
動態(tài)規(guī)劃 + 二次規(guī)劃主流。90 年代，大神提出了時間和空間解耦的數(shù)學(xué)模型，轉(zhuǎn)化為 2 個低維的非線性規(guī)劃問題，降低了搜索難度。
搜索（啟發(fā)式采樣）為主。真車遇到的很多問題，改不動規(guī)劃模型?；诓蓸拥?，有能力修 bug。

個人理解：

當(dāng)問題足夠復(fù)雜時，全局最優(yōu)很難找，投入產(chǎn)出比一般也不高。
產(chǎn)品角度，只要不比人差 & 不比競品差，一般就夠了。一般也不是很依賴全局最優(yōu)。

優(yōu)化問題的經(jīng)典教科書：An Introduction to Optimization
Introduction to Optimization, 4th Edition.pdf
TensorIR 優(yōu)化

在 2-3 層，loop optimization 三板斧:

Fusion
Tiling
vectorization

From https://mlc.ai/summer22/slides/2-TensorProgram.pdf P12 頁開始看

(, 下載次數(shù): 162)
優(yōu)化效果如下圖，F(xiàn)rom https://zhuanlan.zhihu.com/p/613611390
右圖還可以看出，NV 認(rèn)真手動優(yōu)化了 resnet50 和 bert，但 mobilenet 優(yōu)化的少。

(, 下載次數(shù): 169)
MLIR codegen arxiv 總結(jié)不錯，要點：

(, 下載次數(shù): 157)
注：
Async dispatch 就是 split 成多個 block，進行 task parallel 的執(zhí)行。在 GPU，NPU 上很常見。
Blockwise 運算，不是 elementwise 的?；邦悓?Blockwise 運算做 split 時，一般要處理數(shù)據(jù)重疊問題。
Conv 算子的例子：

(, 下載次數(shù): 161)
要點：split 之后，相鄰兩份數(shù)據(jù)有依賴。導(dǎo)致 task parallel 時更復(fù)雜。
深入理解，需要了解 conv 算子的原理，推薦 https://www.youtube.com/watch?v=FmpDIaiMIeA&t=1099s
局限性 & 研究方向

主要的研究難點

動態(tài) shape
Control flow
自動生成新的算子實現(xiàn)。手寫的工作量很大。

其他 AI compiler

現(xiàn)在 ai compiler 可能比 dl 框架還多。這么下去，得做個 compiler for compiler 了
要點總結(jié)：

(, 下載次數(shù): 171)
XLA

XLA 基于 LLVM 框架開發(fā)的，前端的輸入是 Graph，前端沒有將 Graph 直接轉(zhuǎn)化為 LLVM IR
特點：

Theano 之后的第 1 個 ML-compiler。
Compiler 知道每一個 operation 的完整 semantics。operation set 稱為 HLO。影響是：
調(diào)度的實現(xiàn)，簡單了很多。而且，調(diào)度的效果很好。
過度依賴 perfect knowledge of operation semantics。導(dǎo)致代碼復(fù)雜，要把部分 operation 的 semantic 在局部或者全局傳來傳去，比如 fuse pass。
沒有 serializable IR，很難用獨立的 compiler flow 做 UT，inspected，compose。導(dǎo)致代碼的復(fù)用性差，TPU 和 GPU 版本并沒有 share 多少 code。-- XLA 沒把自己當(dāng)一個 compiler 去做。

相比于 TVM 和 MLIR，xla 就是一個早期的探索，

沒啥設(shè)計。就是早期的 shi 山代碼。
XLA 作為早期的 ai compiler 雛形，啟發(fā)和指引意義很大。

Fusion is XLA&#39;s single most important optimization
MLIR

全稱: Multi-Level Intermediate Representation
Chris 繼 LLVM、CLang、Swift 之后第 4 個重磅項目。且有 Google 背書。
The “multi-level” aspect is very important in MLIR: adding new levels of abstraction is intended to be easy and common. Not only this makes it very convenient to model a specific domain, it also opens up a lot of creativity and brings a significant amount of freedom to play with various designs for your compiler: it is actually a lot of fun!
MLIR v.s. TVM
TVM 通過劃分 4 層 IR 的方法，大大的簡化了問題，& 推動了早期的技術(shù)發(fā)展。
但 4 層 IR 之間的信息隔離，導(dǎo)致一些優(yōu)化無法實施，無法跨層搜索 global 的最優(yōu)解。
MLIR 的 multi-level 可以很好的解決這個問題。
個人感覺，這種設(shè)計是雙刃劍。
太過于靈活，放在一般團隊里，如果用不好，過于靈活 == 沒有設(shè)計。
如果是比較優(yōu)秀的團隊，靈活性高則是極好的。
如何入門

通常，大家都會推薦這些書：

編譯原理相關(guān)書籍
- stanford CS143 編譯原理，https://www.bilibili.com/video/BV1NE411376V/
- 鯨書：高級編譯器設(shè)計與實現(xiàn)
- 龍書：編譯原理
- Engineering a Compiler （編譯器設(shè)計）
- SSA-based Compiler Design SSA-based Compiler Design
體系結(jié)構(gòu)相關(guān)
- 計算機體系結(jié)構(gòu) 量化研究方法
- 計算機組成與設(shè)計 RISCV版
編譯相關(guān)經(jīng)典論文
- The Compiler Design Handbook, 2nd Edition (Dec 2007).pdf
- PLDI論文
- ISCA論文
- CGO論文
LLVM 相關(guān)學(xué)習(xí)書籍。low-level IR, optimization, and code-generation
- Getting Started with LLVM Core Libraries （LLVM 編譯器實戰(zhàn)教程）
- LLVM Cookbook
- LLVM Essentials

編譯原理 links
體系結(jié)構(gòu) links
LLVM links
補充說明：

以上材料，我轉(zhuǎn)入這個行業(yè)前都沒讀過。
最近開始讀了。看過跟沒看過，理解問題確實不一樣。
LLVM Cookbook 太貴了，電子版打印下來更靠譜。

更實際的，我的 1 年技術(shù)計劃如下，目前進展 50%，感覺比較實際。

主流 ai compiler
- TVM
- MLIR, XLA
Deep Learning 框架
- Tensorflow, Pytorch
- Onnx, Tensorrt
- Oneflow, paddle
deep learning 模型。
- CV、NLP、搜索、推薦的經(jīng)典模型
- Transformer 深入理解
- 選擇性的看遷移學(xué)習(xí)、強化學(xué)習(xí)、圖模型
c++ 開發(fā)
- 搞的動 TVM 和 tensorrt 涉及的 c++ feature
- 熟練寫多線程、高性能攢 batch 的代碼
傳統(tǒng)編譯器。
- stanford CS143編譯原理
- 刷龍書
- LLVM tutorial
HPC 異構(gòu)編程 & 并行計算。
- GEMM 做一遍
- 體系結(jié)構(gòu), 公司芯片的電路結(jié)構(gòu)
- CUDA 的編程模型

(, 下載次數(shù): 149)

-----------------------------

作者: 林宥嘉 時間: 2023-4-4 10:51
[捂臉]感覺能明白其中一半，就很厲害了

作者: 馬小胖 時間: 2023-4-4 12:42
學(xué)習(xí)目標(biāo)造火箭，實際進度省略一千字

作者: 你牛了 時間: 2023-4-4 14:17
一年竟然能學(xué)完這么多

歡迎光臨愛鋒貝 (http://m.7gfy2te7.cn/)

^{<rt id="gceck"></rt>}^{<kbd id="gceck"></kbd>}