2024龍蜥大會技術實踐分論壇：PyTorchCompile模型優化和應用-朱宏林.pdf

1、Compile 技術概覽、性能優化實踐和應用PyTorch Compile 優化和應用阿里云程序語言與編譯器團隊演講人 朱宏林程序語言與編譯器 我們是阿里云編程語言與編譯器團隊，負責編程語言（C&C+、Java 等）、運行時（PyTorch 等）的研發與產品化，從而支持上層云和 AI 基礎設施并且通過軟硬件結合的方式對系統進行性能優化 我們目前已有產品與核心技術：AJDK(AlibabaJDK)，Alibaba GCC/LLVM，AI 優化軟件棧 等。我們的成果也被等操作系統、軟件工程、計算機領域頂會收錄。團隊介紹概覽PyTorch compile 技術介紹參考資料Compile 流程http

2、s:/pytorch.org/docs/stable/piler_dynamo_overview.html.pyBytecodeCPythonCallbackOriginalModifiedObject Code1.替換默認的 frame evaluation2.第一次執行時進入 callback 觸發編譯3.修改 bytecode 后返回給默認的 frame evaluation 執行Compile 流程經過 Dynamo，AOT Autograd，Inductor 后，生成目標代碼Triton（GPU）C+（CPU）Aten op,cutlass.三種 IRTorch IRAten/Pri

3、m IRLoop-level IR參考資料https:/ IRAten IRLoop-level IRPyTorch compile 實際應用中遇到的問題和解決方案Compile 應用參考資料https:/pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.htmlFused Attention替換成SDPA基于 transformer 結構的模型原始 self-attention 實現優化實現，調用 SDPA，flash-attn 等SDPAFlashAttentionMemory-Efficien

4、t AttentionPyTorch implementationPattern Matching參考資料Attention with Linear Biases paper:https:/arxiv.org/pdf/2108.12409沒有使用 SDPA使用 SDPA，但進入了較慢的 backend性能限制：Q K.T*+*ALiBimask+softmax()V合并 opSDPA(,mask,)memory efficient attn原始論文圖Constant Folding性能限制：沒有使用 SDPA使用 SDPA，但進入了較慢的 backendA0+AAA復雜圖=0+AAQ K.T*

5、mask+softmax()VSDPAMemory-EfficientFlashwith maskpattern matchingmask 矩陣全0Constant Folding說明：測試 GPU A10，數據類型 bfloat16，吞吐表示每秒執行多少次 self-attention 計算vanilla,memory efficient 和 flash attention 的性能差別Llama2-7B(1,2048,32,128)XLM-RoBERTa(1,8192,16,64)優化場景 確保 dynamic shape 下能夠進行sdpa 的匹配。增強 constant folding 能力，去除不需要的 mask，從而使得 sdpa 進入更快的 flash 路徑。XLM-RoBERTa(1,8192,16,64)PyTorch：編譯優化、算子優化、運行時優化 三方庫：deepspeed、vllm、triton、oneDNN/ACL 等 CPython：GC 優化應用輸出研發范圍AC2（Alibaba Cloud AI Containers）鏡像 產品文檔：https:/ ac2-