OpenXLA Project

Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

`-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`

プログラムからメッシュオペレーションとシャーディング表記を削除します。

オプション

-keep-module-sharding-metadata : Whether to keep module-level sharding metadata (sdy.parameters_shardings and sdy.output_shardings).

`-sdy-drop-sharding-rules`

登録されているすべてのオペレーションから OpShardingRuleAttr を削除します。

`-sdy-export-named-computations`

NamedComputationOp から呼び出しをアウトライン化します。

NamedComputationOp を、NamedComputationOp の name という名前の新しいプライベート関数を持つ CallOp に変換するパスを作成します。新しい FuncOp と CallOp は、元の NamedComputationOp のオペランド/結果と同じシャーディングを持ちます。

モジュールに NamedComputationOp と同じ名前の関数がある場合、MLIR シンボルテーブルはそれを {name}_# に変更します。

`-sdy-insert-explicit-reshards`

明示的なリシャーディングを挿入して、すべてのオペレーションが互換性のあるシャーディングを持つようにします。

互換性のあるシャーディングとは、基本的に、オペレーションがシャーディングされたオペランドを受け入れ、リシャーディング通信を必要とせずにシャーディングされた結果を生成できることを意味します（オペレーションでは、all-reduce や halo-swaps などの通信が必要になる場合があります）。

伝播後も、一部のオペレーションでは互換性のないシャーディングが残る可能性があります。

軸（またはサブ軸）を使用して、複数のテンソル間で対応しないディメンション（matmul の非縮約ディメンションなど）をシャーディングする場合や、軸が 1 つのテンソルのディメンションをシャーディングするが、もう 1 つのテンソルの対応するディメンションをシャーディングしない場合、オペレーションにシャーディングの競合があると見なされます。したがって、このパスの後、オペレーションは競合のない状態になります。

このパスは、リシャーディングオペレーションを明示的に挿入します。これにより、各オペレーションで、対応するディメンションがすべてのオペランドと結果で同じようにシャーディングされ、各軸（またはサブ軸）は単一のディメンションタイプのみをシャーディングするために使用できます。

例:

入力:

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.
-mark-partial-result-with-unreduced-axes : Whether to mark partial reduction results with unreduced axes.

`-sdy-insert-func-call-reshards`

func と call のシャーディングの競合に対してリシャーディングを挿入します。

結果に対する func と call のシャーディングの競合に対してリシャーディングを挿入します。

`-sdy-optimize-collectives`

集合通信オペレーションを最適化します。

分割されたディメンションを分解し、結合された all-to-all に融合することで、all-to-all チェーンの前に冗長な collective-permute オペレーションを排除します。

`-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, func=subroutine'.
-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).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

`-sdy-propagate-to-func-results`

func ターミネータから func の結果にシャーディングを伝播します。

func ターミネータ値のシャーディングを対応する func.func の結果にコピーし、メイン関数の既存のシャーディングを保持します。

`-sdy-remove-all-gather-reduce-scatter-for-cmv1`

CMV1 の sdy.all_gather と sdy.reducescatter を削除します。

パターン all-gather + dot の all-gather を削除します。パターン dot + reduce-scatter の reduce-scatter を削除します。このパスは、collective 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, &q>uot;y"=<2, ">;z"=2
%0 <: tensor16x2xf32 {sdy.sharding>@mesh, \[{"x",< "y", &quo>t;z"<}, {}\]<>/span>
%1 = sdy.reshard %arg0 @mesh, \[{"x"}, {}\] : tensor16x2xf32

出力:

mesh = <"x"=2, &q>uot;y"=<2, ">;z"=2
%0 <: tensor16x2xf32 {sdy.sharding>@mesh, \[{"x", "y", "z"}, {}\]<
%1 = sdy.all_gather> \[{&quo<t;y">;, "z"}, {}\] %arg0 out_sharding=@mesh, \[{"x"}, {}\] : tensor16x2xf32

上記の例では、テンソル %0 : tensor<16x2xf32> は \[{"x", "y", "z"}, {}\] としてシャーディングされます。次に、reshard オペレーションで \[{"x"}, {}\] としてリシャーディングされます。最初軸では、リシャーディング後に接尾辞 {"y", "z"} が削除されるため、 {"y", "z"} が all-gathered されたと推測できます。2 番目のディメンションは変更されません。

オプション

-keep-redundant-reshards : Whether it keeps redundant reshards or removes.

`-sdy-resolve-permutation-factors`

kPermutation ファクタにマッピングされたディメンションのシャーディングを解決します。

kPermutation ファクタでディメンションをシャーディングするには、デバイス間通信が必要になる場合があります（ウィンドウの halo 交換や、逆方向の集合置換など）。

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).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

`-sdy-resolve-single-device-sharding`

ターゲットデバイス ID によって保護されている単一デバイスオペレーションを stablehlo に変換します。

オプション

-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 などの境界オペレーションであるか）を確認します。