（本文閱讀時間：6分鐘）
目前，深度學習已經透過多種應用形式廣泛深入人們的日常生活當中。但隨著模型規模的不斷增大，其訓練變得耗時且昂貴。為此，科研人員不斷設計出各種並行策略來提高其在多裝置上的執行效能。
在這個過程中，來自微軟亞洲研究院的研究員們注意到現有的方法依然會遺漏許多潛在的高效並行策略。如何充分探索並行策略的可能性從而進一步釋放效能，成為研究員們關注的焦點。
為了解決現存的並行策略問題，研究員們重新審視了並行策略的現有實現方式。目前主流的訓練系統（如 Megatron-LM、DeepSpeed、Alpa 等）中通常內建了數種並行策略，例如，資料並行（data-parallelism）、張量並行（tensor-parallelism）和流水線並行（pipeline-parallelism）等，並且允許透過配置實現策略的聯合應用。這樣的方式對於使用者來說較為簡便，但仍然可能遺漏高效的策略。而且當人們需要使用新的高效並行策略時，也會因為只能透過修改大量的系統底層程式碼來實現而顯得捉襟見肘。
為此，研究員們重新思考了並行策略的基本構成。眾所周知，深度學習的模型執行通常可以表示為一張資料流圖（data-flow graph），其中以表示資料的張量（tensor）作為圖邊，張量的計算運算元（operator）作為頂點。並行化的過程實際上就是將原本針對單一裝置執行的計算資料流圖重新進行分割變換，從而變成針對多個裝置執行的分佈化資料流圖。因此，研究員們提出了一套基本操作作為並行化原語，包含：
1） op-trans 表示如何對運算元和張量進行劃分；
2） op-assign 表示對劃分後的運算元設定哪臺裝置來執行；
3） op-order 可以設定同一臺裝置上運算元的執行先後順序。
利用這套原語，資料流圖中的每個運算元和張量如何劃分，以及在空間（裝置間）和時間（同一裝置）上如何排程，都能夠得到準確的描述和解釋。
圖1：深度學習模型的資料流圖及分佈化中的“時間-空間”排程
並行原語可以用於表示目前廣泛應用的各種並行策略。例如，資料並行就可以表達為將所有前向（forward pass）和後向（backward pass）運算元進行資料樣本維度的劃分，均勻地分配到所有的裝置上，而將最佳化器（optimizer）的運算元複製到每臺裝置上。並且所有裝置上的運算元保持原圖中相同的執行順序。透過這樣一種並行原語的引入，研究員們就能使各種並行策略在同一個框架下進行描述和融合，從而極大地拓展了並行策略空間的表示範圍。
這種通用性不僅系統化地表達了現有的並行策略，還為探索新的並行策略提供了可能。例如，對於產生單裝置無法容納中間結果的運算元，原有的策略自然是在張量並行的劃分後，分配到多個裝置上並且採用通訊協調，而透過原語僅對該運算元進行張量劃分後，即可全部排程到同一個裝置上順序執行，成功執行的同時也避免了可能產生的通訊開銷。
圖2：引入並行原語對並行策略進行表示
儘管並行原語擴充套件了策略空間，但這種擴充套件也帶來了新的挑戰——巨大的策略空間讓搜尋在有限的時間內完成變得非常困難。面對海量的可能性組合，如何高效地找到最佳策略成為了一個亟待解決的問題。
研究員們的解決方案是藉助“專家”的經驗智慧，進行有效的策略搜尋限定。這種指導透過並行原語進行描述，可以無縫銜接到研究員們的系統中。例如，下圖中研究員們透過 op-trans 的引數限定運算元 op 只使用 algo 集合指定的劃分方案，並且只切分為 num 限定的份數。
圖3：並行原語表示策略搜尋限定
“專家”透過設定特定的約束條件，能夠大幅度縮小搜尋空間，使得搜尋過程更加高效且有針對性。在實驗中，研究員們觀察到透過約束搜尋空間可以獲得高達10倍的搜尋效率提升，同時不損失策略的效果。
利用這種方法，研究員們不僅發現了現有方法所遺漏的高效能策略，還能夠在較短的時間內完成搜尋，提升深度學習訓練的效率。目前這一系統已經在多個深度學習模型的訓練中得到了驗證，展示了顯著的效能提升。
更多細節可參考論文詳情：
nnScaler: Constraint-Guided Parallelization Plan Generation for Deep Learning Training
https://www.usenix.org/conference/osdi24/presentation/lin-zhiqi
藉助引入並行原語和專家智慧指導的策略搜尋，微軟亞洲研究院的研究員們有效地解決了深度學習並行策略設計中的諸多問題，並在極大地拓展深度學習並行策略空間的同時，為未來的並行策略研究提供了新的方向和工具。研究員們非常期待這一方法能夠在更廣泛的應用中展示其潛力，為深度學習的發展帶來更多的可能性。
你也許還想看：
