本文分享自華為雲社群《深度解讀昇騰CANN模型下沉技術，提升模型排程效能》，作者：昇騰CANN。
AI模型的執行通常情況下需要CPU和NPU（昇騰AI處理器）等AI專用處理器協同工作，CPU所在位置稱為主機端（Host），而NPU所在位置稱為裝置端（Device）。對於採用Host排程的AI模型來說，Host下發Task的時序和Device執行Task的時序是非同步的，如果Device執行Task的速度比Host下發Task的速度快，則Device會處於空閒狀態。比如，大模型場景的增量推理或訓練的FineTune階段，都是計算量較小的場景，此時很容易出現單個運算元的Host下發時間比Device上的運算元執行時間還長，從而導致Device間歇處於空閒狀態。這種現象通常稱為Host Bound，這種模型也稱為Host Bound模型。
如何減少Host Bound模型的Device空閒時間，從而最佳化模型執行效能顯得尤其重要，GE（Graph Engine）圖引擎透過圖模式的Host排程和模型下沉排程的方式，可提升模型排程效能，縮短模型E2E執行時間。
1 模型Host排程
Host CPU將模型中的運算元依次下發到Device執行（如下圖中的標號①所示），每一個運算元在執行流上以1個或多個Task的形式存在，昇騰AI處理器依次拉取執行流上的Task執行（如下圖中的標號②所示），我們稱這個過程為Host排程。Host排程需要Host和Device頻繁互動，在實際的訓練或推理場景中，模型會執行多次，每次執行都會觸發Host把模型上的所有運算元遍歷下發一遍。
圖1 Host排程：
1.1 Device Bound與Host Bound
如果Device上Task執行速度比Host下發慢，則Device上Task排程開銷和運算元執行時間成為瓶頸，這種模型通常稱為Device Bound模型，如圖2所示。Device Bound的模型，Task的下發開銷會被已下發Task的執行時間掩蓋起來。
圖2 Host排程場景Device Bound執行時序分析：
如果Device上Task執行速度比Host下發快，則Host排程開銷成為瓶頸，這種模型通常被稱為Host Bound模型，如圖3所示。Host Bound的模型，Task的下發開銷不能完全被已下發Task的執行時間掩蓋，Device執行時序上存在間歇的空閒狀態。
圖3 Host排程場景Host Bound執行時序分析：
1.2 單運算元模式Host排程與圖模式Host排程
在前面幾期介紹中，我們知道當前業界主流的深度學習框架提供了Eager執行（Eager Execution，即時執行）模式與圖執行模式。Eager模式也叫單運算元模式，它是一種Host排程模式，由Host CPU逐個下發運算元，一個運算元的下發流程包含Python處理、Python到C++資料結構轉換、Tiling計算，申請運算元的Workspace記憶體和輸出記憶體，Launch等Host操作。為了加速單運算元Host排程，在PyTorch中，昇騰適配層採用了生產者—消費者雙執行緒模式加速，生產者執行緒主要負責Launch之前的處理動作，消費者執行緒主要負責Launch運算元。
相比於單運算元模式，圖模式的Host排程可以避免總是返回Python呼叫棧，避免冗餘流程與資料結構轉換，並且可以直接使用圖編譯階段完成的Infer Shape與Tiling計算結果。因此，圖模式Host單執行緒排程與單運算元模式Host雙執行緒排程相比，排程效能持平或略優。
2 模型下沉排程
圖模式排程可以降低Host側的排程耗時，並一定程度減少模型執行E2E耗時，但如何降低Device上執行時序的空閒時間仍是需要考慮的問題。對於靜態shape模型，昇騰支援下沉排程，可大幅最佳化Host排程效能。
所謂靜態shape模型，是指模型每次執行的輸入tensor shape是固定不變的，模型中的所有運算元，給定輸入shape後，都可以確定自己的輸出shape，那麼該模型為靜態shape模型。
靜態shape模型在編譯時即可確定所有運算元的輸入輸出shape，結合昇騰記憶體複用演算法，可完成模型級記憶體編排；靜態shape模型在編譯時還可提前完成所有運算元的Tiling計算等Host側計算。綜上，靜態shape模型可以在編譯時完成所有執行時的Host排程工作，這是靜態shape模型下沉排程的理論基礎。
2.1 下沉排程原理
模型下沉排程分為兩個階段，模型載入和模型執行。
圖4 模型下沉示意圖：
- 模型載入
模型載入的具體動作和Host排程類似，即遍歷圖中的所有運算元並將其整體下發至Device流上，區別在於下發到流上不立即執行。模型載入是一次性的動作，在首次模型執行時完成模型載入，如圖4中的過程①所示。
- 模型執行
模型載入完成之後，可以像下發單運算元Task一樣，向執行流下發一個模型執行Task，昇騰AI處理器排程到該Task時（如圖4 “執行流”中的E），執行模型中所有Task（如圖4中的過程③）。如果需要多次執行模型，僅需多次下發模型執行Task，如圖4中的過程②所示。
Host Bound排程和模型下沉排程的時序比較如圖5所示，可以看出模型下沉執行的開始有一個模型下發的頭開銷，模型下沉執行E2E會有一個相對於Host排程的收益，模型的下沉頭開銷越小，收益將越大。
圖5 模型下沉排程Host/Device時序分析：
每次模型下發時，支援更新模型的Feature Map記憶體地址和輸入輸出記憶體地址，如果模型的Feature Map記憶體和模型輸入輸出記憶體發生了更新，則在模型下沉頭開銷（即上圖中的m_l_t，後文有介紹）中會完成模型內運算元相關地址的重新整理。
2.2 下沉效果
模型下沉執行存在如下優勢：
- 減少CPU的負載
對於模型下沉執行的方式，執行時不需要Host和Device頻繁互動，排程的開銷僅包含模型下沉的頭開銷和Task從流上排程到加速器上的開銷。總的來說，模型下沉執行的方式，Host CPU的負載降低了，在一定程度上整機、叢集的功耗也會降低。
- 減少Rank之間的通訊抖動
對於Host排程執行方式，叢集場景下由於不同卡上的Host下發時序不一定完全同步，卡間集合通訊可能存在一定程度上的抖動。而下沉執行方式則不存在該問題，因為Task已提前排布到Device上。
- 提升E2E的收益
由於Task已提前下沉到了Device，對於Host Bound的模型，E2E會有效能提升。以Host Bound的LLaMA-7B模型的推理場景為例，圖6是Host排程的Profiling效能資料，圖7是模型下沉執行的Profiling效能資料。透過Profiling結果可以看出該模型是一個Host Bound模型，採用Host排程方式，Device上計算時間和空隙的時間比例接近1:1，採用模型下沉執行方式，Device空隙大幅減少，並且端到端有18ms的收益，吞吐量整體提升37%。
圖6 LLaMA-7B Decoding Host排程Profiling：
圖7 LLaMA-7B Decoding模型下沉Profiling：
2.3 下沉頭開銷
下沉執行頭開銷包括以下幾部分：
1、模型輸入輸出Tensor到內部InputData/OutputData資料結構轉換。如圖8中的stage1_t。
2、如果模型的Feature Map記憶體和模型輸入輸出記憶體有變更，重新整理相關聯運算元的相關地址。如圖8中的stage2_t。
3、如果模型的輸入不支援零複製，（如模型的輸入為Host記憶體），則下發非同步複製Task。如圖8中的stage3_t。
4、下發模型執行Task。如圖8中的stage4_t。
圖8 模型下沉頭開銷拆分：
以盤古71B增量推理模型為例（模型輸入輸出總個數約1600個，模型內節點總數約6300個，模型執行時，Feature Map記憶體地址不變更，10個輸入輸出記憶體地址變更），當前模型下沉的頭開銷約2ms，關於頭開銷的效能最佳化在持續進行中。
3 更多介紹
GE模型下沉技術的相關介紹就到這裡，歡迎大家關注後續技術分享。如需獲取更多學習資源請登入昇騰社群。
點選關注，第一時間瞭解華為雲新鮮技術~