Shardy 是適用於所有方言的 MLIR 向量分割系統。這項工具是 GSPMD 和 PartIR 團隊合作開發，結合了兩個系統的優點，以及兩個團隊和使用者共同分享的經驗。
優點
- 結合 GSPMD 的傳播與 PartIR 的漸進式分區，為使用者提供更多控制權和可預測性。
- 由共用體驗驅動的全新功能，例如重新調整形狀的創新支援功能，這類功能經常會產生額外的通訊，除非使用者知道如何處理，否則會造成困擾。
- 改善可用性和偵錯能力，以提高使用者速度，例如使用以軸為基礎的分割表示法。
- 使用 MLIR 的簡單開放原始碼程式碼庫，並有更廣泛的活躍貢獻者 (內部、外部和跨時區) 來支援使用者。
元件
- 分割表示法：以軸為基礎的分割表示法，會繫結至特定邏輯網格 (可能來自多個網格)，並支援限制維度分割和軸、分割軸以便執行重塑等作業、優先處理遞增分割作業等等。
- 編譯器 API：一組可與區隔表示法一併使用，用於影響區隔傳播的編譯器元件。
- 輸入/輸出分割 - 將分割附加至主函式的輸入或輸出內容，以指出在函式傳遞/傳回時，輸入/輸出張量應如何分割。
- 區隔限制：將區隔附加至中介張量 (例如 matmul 的結果)，以表示應如何區隔該張量或其用途的子集。
- Shard As/Like：依據 ID 將多個張量分組，表示應以相同方式分割。
- 手動運算：包含使用網格軸子集手動分割的子運算，其中會為所有輸入和輸出指定沿著這些手動軸子進行分割，並在子運算內部將張量類型設為相對於這些分割的本機。
- 分割區傳播：結合使用者優先順序和分割區限制的傳播演算法，搭配編譯器成本模型和推論法：
- 使用者定義的優先順序，例如先執行批次平行處理，再執行 ZeRO
- 以運算為基礎的優先順序，例如先以元素為單位的運算，再以矩陣相乘為單位的運算等。
- 更精細的規則，例如偏好批次維度。
- SPMD 區塊劃分器：這個元件會將程式劃分為 SPMD 程式，並在程序中加入必要的資料移動/格式設定和集體作業，以降低分割傳播決策。
- 短期內，初始實作會使用目前的 GSPMD SPMD 分割器。
- 長期來說，我們預計建立新的 MLIR 版 SPMD 分割器。
程式碼存放區
Shardy 專案目前正在積極開發中，我們希望能向開放原始碼社群徵求意見回饋。您可以在 https://github.com/openxla/shardy 取得 Shardy 程式碼