-sdy-close-shardings
關閉張量分片,並捨棄複製的軸。
-sdy-constant-or-scalar-merger
合併相同的常數和純量擴展,並比對分片。
對分片相同的常數執行輕量 CSE。
匯入管道會分割及複製常數和純量擴展,因此分片不會在常數子運算的各種用途之間傳播。如果常數在傳播後具有相同的分片,這個階段會合併常數,以節省編譯時間。詳情請參閱 -sdy-constant-or-scalar-splitter。
-sdy-convert-global-to-local
將 SDY 程式從全域形狀轉換為本機形狀。
根據分片屬性分割邏輯維度,將 SDY 程式從全域形狀轉換為本機形狀。
這個傳遞會運用型別轉換器,將全域邏輯形狀的 RankedTensorType 對應至裝置本機的實體形狀。
選項
-per-dim-all-gather : Keep per-dimension all-gather without combining them into a single all-gather.
-combine-multi-dimension-reduce-scatter : Combine multi-dimension reduce-scatter into a single reduce-scatter.
-enable-rgv3 : Use StableHLO ReplicaGroupV3 (mesh-axes based) for collectives.
-replica-count : Number of replicas (data parallelism).
-partition-count : Number of partitions (model parallelism).
-sdy-drop-sharding-and-mesh
從程式中移除網格作業和分片註解。
-sdy-drop-sharding-rules
所有已註冊的藝人OpShardingRuleAttr的開賣活動。
-sdy-export-named-computations
列出 NamedComputationOp 的通話內容。
建立傳遞項目,將 NamedComputationOp 轉換為 CallOp,並使用名為 NamedComputationOp 的新私有函式 name。新的 FuncOp 和 CallOp 與原始 NamedComputationOp 的運算元/結果具有相同的分片。
如果模組中有與 NamedComputationOp 同名的函式,MLIR 符號表會將其變更為 {name}_#。
-sdy-insert-explicit-reshards
插入明確的重新分片,使所有作業都有相容的分片。
相容的分片基本上是指作業可以接受分片的運算元,並產生分片的結果,不需要任何重新分片通訊 (請注意,作業可能仍需要通訊,例如全縮減或光環交換)。
傳播後,部分作業可能仍有不相容的分片。
請注意,當軸 (或子軸) 用於跨多個張量分片非對應維度 (例如 matmul 中的非收縮維度),或當軸在一個張量中分片維度,但不在另一個張量中分片對應維度時,表示運算發生分片衝突。因此,在這次傳遞後,作業就會變成無衝突。
這個傳遞會明確插入重新分片作業,以便針對每個作業,在所有運算元和結果中,以相同方式分片對應的維度,且每個軸 (或子軸) 只能用於分片單一維度類型。
範例:
輸入:
mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
: (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
輸出內容:
sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
: (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
在上例中,lhs 和 rhs 都在軸「x」上,以非收縮維度進行分片,因此不相容。傳遞會在點運算之前,在 rhs 上插入明確的重新分片,以便點運算具有相容的分片。
選項
-enable-full-version : Enable full version.
-sdy-insert-func-call-reshards
插入 func 和呼叫分片衝突的重新分片。
插入 func 和呼叫分片,以解決結果的分片衝突。
-sdy-optimize-collectives
最佳化集體通訊作業。
分解分割維度並融合為合併的 all-to-all,藉此消除 all-to-all 鏈結前多餘的集體排列作業。
-sdy-pad-for-divisibility
將無法整除的分片張量填補為可整除的形狀。
-sdy-per-instruction-partitioning
_選擇性擷取目標指令,對這些指令執行分割器管道,產生裝置本機程式碼,並將這些程式碼包裝在 sdy.manualcomputation 中。
為進行偵錯和二分作業,請使用獨立的 Shardy 分區器管道 (resolve-permutation-factors、reshard-to-collectives、pad-for-divisibility、convert-global-to-local),選擇性地分割個別指令,並將產生的裝置本機程式碼包裝在 sdy.manual_computation 中。這樣可確保全域程式的其餘部分保持不變,因此 XLA SPMD 只會分割未包裝的全域指令。
選項
-filter : Filter string for selective partitioning. Can be empty (all sharded ops), comma-separated op name substrings (e.g. 'dot, pad'), or key-values like 'selectLow=0, selectHigh=10'.
-enable-halo-exchange : Implement halo exchange logic for windowed operations inside the sub-pipeline.
-replica-count : Number of replicas (data parallelism).
-partition-count : Number of partitions (model parallelism).
-sdy-propagate-to-func-results
將分片從函式終止符傳播至函式結果。
將函式終止符值的分片複製到對應的 func.func 結果,保留主要函式的現有分片。
-sdy-remove-all-gather-reduce-scatter-for-cmv1
_Removes sdy.all_gather and sdy.reducescatter for CMV1.
移除模式 all-gather + 點中的 all-gather。移除模式點 + reduce-scatter 中的 reduce-scatter。這個傳遞是為了與集合 matmul V1 (CMV1) 相容。這是 b/432019089 的暫時解決方案。
-sdy-remove-propagation-debug-info
匯出時移除傳播偵錯資訊 (傳播邊緣和來源分片)。
-sdy-remove-sharding-groups
在傳播後移除 ShardingGroupOps。
-sdy-remove-sub-axes-in-input-output-shardings
移除輸入/輸出分片中的子軸。
部分 Shardy 使用者希望函式輸入/輸出內容具有分片,但沒有子軸。這個傳遞會從輸入/輸出開放維度分片中移除子軸及其後方軸。這個階段通常會在 sdy-update-non-divisible-input-output-shardings 之後,確保移除子軸不會導致任何無法分割的分片。
-sdy-reshard-to-collectives
將 ReshardOp 轉換為各種 Shardy 集合運算。
比對重新分片作業,並將其重新編寫為各種 Shardy 集合作業。通過此階段後,模組中就不會再有任何重新分片作業。
如果 keepRedundantReshards 為 true,則只會保留多餘的重新分片作業。根據預設,系統會假設已插入明確的重新分片 (sdy-insert-explicit-reshards),且不會保留多餘的重新分片。如果可能尚未插入明確的重新分片,則應保留多餘的重新分片。
範例:
輸入:
mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>
輸出內容:
mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.all_gather \[{"y", "z"}, {}\] %arg0 out_sharding=<@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>
在上述範例中,張量 %0 : tensor<16x2xf32> 會分片為 \[{"x", "y", "z"}, {}\]。接著,會以 reshard op 重新分片,如 \[{"x"}, {}\]。在第一個軸上,由於後綴 {"y", "z"} 會在重新分片後移除,因此我們推斷已全數收集 {"y", "z"}。第二個維度不會變更。
選項
-keep-redundant-reshards : Whether it keeps redundant reshards or removes.
-sdy-resolve-permutation-factors
解決對應至 kPermutation 因數的維度分片問題。
使用 kPermutation 因子分割維度可能需要跨裝置通訊 (例如視窗的暈輪交換或反向的集體排列)。
如果 enableHaloExchange 為 true,傳遞內容會使用可用的最佳化通訊邏輯來解析排列組合因素。否則,傳遞會單純插入 sdy.reshard 作業,以複製這些維度。enableHaloExchange 的預設值為 true。
選項
-enable-halo-exchange : Implement halo exchange logic for windowed operations.
-replica-count : Number of replicas (data parallelism).
-partition-count : Number of partitions (model parallelism).
-sdy-resolve-single-device-sharding
將單一裝置作業降低至 stablehlo,並以目標裝置 ID 做為防護。
選項
-replica-count : Number of replicas per partition.
-partition-count : Number of partitions per replica.
-sdy-sharding-constraint-to-reshard
將 ShardingConstraintOp 轉換為 ReshardOp。
-sdy-sink-data-flow-edges
將所有 DataFlowEdgeOp 匯入輸入內容。
將每個 DataFlowEdgeOp 的分片移至其輸入內容 (邊緣的根目標),並以輸入內容取代作業。
選項
-sink-debug-sharding-origins : Whether to sink the debug sharding origins info. See `debug-sharding-origins` option in propagation for more info.
-sink-debug-propagation-edge-sharding : Whether to sink the debug propagation edge sharding info. See `debug-propagation-edge-sharding` option in propagation for more info.
-sdy-sink-func-data-flow-edges
將所有 FuncDataFlowEdgeOp 匯入輸入內容。
將每個 FuncDataFlowEdgeOp 的分片作業移至輸入內容,並以輸入內容取代作業。
-sdy-unflatten-call-graph
取消扁平化呼叫圖。
取消扁平化圖表。這個函式會根據函式附加的「original_func_name」屬性,移除具有相同輸入/輸出分片 和相同來源的函式。
選項
-dedup-functions-fully : If true, regardless of the input and output shardings of functions, it keeps one callee function for each caller function. The default is false, meaning it will deduplicate only if the input and output shardings are the same.
-sdy-update-non-divisible-input-output-shardings
平均分配 FuncOp 輸入/輸出內容,避免因無法整除的分片而需要填補。
Shardy 使用者希望函式輸入/輸出可平均分割/分片,避免需要填補張量。傳播可能會導致輸入/輸出具有不可分割的分片,因此這個傳遞會將其更新為原始分片的最大維度分片前置字串,該前置字串會均勻分片。
-sdy-verify-unreduced-axes
驗證未縮減軸用量的一致性。
驗證每個作業,如果運算元有未縮減的軸,作業會明確縮減這些軸 (例如透過 sdy.reshard),或將這些軸傳遞至結果 (或是 func.call 等界限作業)。