方言に依存しないシャーディング

長期的な目標は、Shardy を完全にスタンドアロンのコンポーネントにして、任意の MLIR 方言で動作できるようにすることです。現在、Shardy は StableHLO に直接依存していますが、Shardy をより柔軟にするために、さまざまな抽象化とインターフェースを通じてその依存関係を解除する作業が進んでいます。

シャーディング ルール

シャーディング ルールは、オペレーションを伝播する方法をエンコードします。Shardy は StableHLO に依存しているため、stablehlo オペレーションごとにシャーディング ルールを定義します。さらに、Shardy には ShardingRuleOpInterface が用意されており、オペレーションで方言オーナーが独自のオペレーションのシャーディング ルールを定義できます。オペレーションがこのインターフェースを実装している限り、Shardy はそれを介して伝播できます。

def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
  let methods = [
    InterfaceMethod<
      /*desc=*/[{
        Returns the sharding rule of the op.
      }],
      /*retType=*/"mlir::sdy::OpShardingRuleAttr",
      /*methodName=*/"getShardingRule"
    >,
  ];
}

データフロー運用

リージョンベースのオペレーションなど、一部のオペレーションでは、すべてのオペランドと結果のディメンション間の対応のみを記述するシャーディング ルールでは不十分な別のアプローチが必要です。このような場合、Shardy は ShardableDataFlowOpInterface を定義します。これにより、方言オーナーはオペレーションを介してシャーディングのプロパゲーションを記述できます。このインターフェースには、オーナーを介して各データフロー エッジのソースとターゲットを取得するメソッドと、エッジ オーナーのシャーディングを取得して設定するメソッドが用意されています。

def ShardableDataFlowOpInterface :
    OpInterface<"ShardableDataFlowOpInterface"> {
  (get|set)BlockArgumentEdgeOwnerShardings;
  (get|set)OpResultEdgeOwnerShardings;
  getBlockArgumentEdgeOwners;
  getOpResultEdgeOwners;
  getEdgeSources;
  // ...
}

データフロー オペレーションの処理方法の概要については、データフロー オペレーションもご覧ください。

インターフェースがまだ実装されていない

今後、Shardy をより柔軟で方言に依存しないものにするために、さらに多くのインターフェースとトレイトが追加される予定です。以下にその一部を記載いたします。

定数分割

MLIR のほとんどのテンソル プログラムには、その値を必要とするオペレーションによって再利用される定数のインスタンスが 1 つあります。これは、必要な定数が同じ場合に意味があります。ただし、プログラムの最適なシャーディングのために、定数の各使用に独自のシャーディングを許可し、他のオペレーションがその定数を使用する方法の影響を受けないようにする必要があります。

たとえば、次の図では、add がシャーディングされている場合、divide と subtract（計算の異なる部分）がシャーディングされる方法に影響しません。

定数分割

これは偽の依存関係と呼ばれます。定数は低コストであるため、同じ定数を使用するオペレーション間に実際の依存関係はありません。そのため、ユーザーは定数（および定数に似た）オペレーションのシャーディングを決定できます。その定数の各使用には、定数のサブ計算の独自のコピーに個別に伝播できる異なるシャーディングを設定できます。

これを実現するには、Shardy ユーザーは次のことを定義する必要があります。- your_dialect.constant -> sdy.constant パス。- sdy::ConstantLike トレイト（iota など）。- 要素ごとのオペレーション（add、multiply など）用の mlir::Elementwise トレイト。- オペレーション（slice/broadcast など）用の sdy::ConstantFoldable。これらのオペレーションは、すべてのオペランドまたは結果が定数である場合、技術的にはコンパイル時に計算できます。

オペレーションの優先度

GSPMD では、要素ごとのオペレーションが最初に伝播され、次に matmul などのオペレーションが伝播されます。Shardy では、方言を事前に把握できないため、ユーザーが独自の OP 優先度を設定できるようにしています。そのため、Shardy に伝播する順序でオペレーションのリストを渡すように依頼します。

次の図は、GSPMD で優先度を使用してオペレーションを適切な順序で伝播する方法を示しています。

オペレーションの優先度。オペレーションの優先度が重要な理由については、GSPMD の論文をご覧ください。

オペレーションの優先度が重要な理由については、GSPMD の論文をご覧ください。

方言に依存しない

上記のインターフェース、トレイト、パスを実装していれば、Shardy は方言で動作します。Shardy をより柔軟で方言に依存しないものにする取り組みを進めておりますので、今後の最新情報にご期待ください。