エラーコード: E2001

カテゴリ: コンパイル時: ハードウェアでサポートされていない RHS DataType

このエラーは、行列乗算（jax.lax.dot_general、jax.lax.conv、jax.numpy.matmul、@ 演算子など）の右辺（RHS）オペランドに使用されるデータ型が、使用されている特定の TPU 世代でネイティブにサポートされていない場合に発生します。

エラー メッセージの例:

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

XLA バックエンド: TPU

概要

TPU の Matrix Multiply Unit（MXU）は、すべてのハードウェア世代で Float32 演算をネイティブにサポートしています。

ただし、BFloat16 や他の量子化データ型（Int4、Int8、Float8 など）のネイティブ サポートは、ハードウェア世代によって異なります。このエラーは、特定の TPU 世代に実行するための物理回路がないデータ型を使用して、カーネルが行列乗算を MXU にマッピングしようとしたときに発生します。

このエラーは通常、コンパイラの Canonicalization パス（サポートされていない型をサポートされている型に自動的に変換しようとするパス。ソフトウェア エミュレーションなどを使用）が有効な変換ルールを見つけられなかったか、互換モードが無効になっているために変換ルールを見つけられなかったことを示します。

デバッグ

このエラーを解決するには、データ型をハードウェアの機能に合わせる必要があります。次のいずれかをご検討ください。

1. ネイティブ型にキャストする

最も確実な修正方法は、matmul オペレーションの前に、カーネル内でオペランドをハードウェアでサポートされているデータ型（TPU v4+ の Float32 や BFloat16 など）に手動でキャストすることです。

  • 理由: Float32 は、すべての TPU 世代の MXU でネイティブにサポートされているユニバーサル データ型です。
  • トレードオフ: VPU（ベクトル処理ユニット）のコスト（キャストの実行に必要なサイクル）が発生しますが、カーネルが現在のハードウェアで実行されることが保証されます。

2. 互換モードを確認する

通常、コンパイラはデフォルトで有効になっている互換モードで、これらの型の不一致の問題を自動的に処理できます。XLA 構成を再確認して、--xla_mosaic_compat_mode が false に設定されていないことを確認します。

これは「ポリフィル」として機能し、ハードウェアがネイティブにサポートしていないオペレーションのソフトウェア エミュレーション シーケンスを挿入します。

互換モードで有効になる機能:

  • 混合精度 MatMul: キャスト オペレーション（matmul の前に整数を Float32 に拡張するなど）を自動的に挿入することで、整数オペランドと浮動小数点アキュムレータを混合できます。
  • 低精度のエミュレーション: 特定のハードウェア世代では、実行前に 4-bit 浮動小数点（4E2M1FN）や 8-bit 浮動小数点（8E4M3FN）などのサポートされていない型を BFloat16 や Float32 などのサポートされている型に拡張して、エミュレートします。

このモードでは、エミュレーションで MXU がデータを処理する前にデータ形式を変換するための追加の命令が必要になるため、ピーク パフォーマンスよりも互換性が優先されます。

3. ハードウェアをアップグレードする、またはサポートをリクエストする

アルゴリズムで、キャストやエミュレーションのオーバーヘッドなしで Int4 や Float8 などの型にネイティブ パフォーマンスが厳密に必要となる場合は、ネイティブ サポートを備えた新しい世代の TPU で実行する必要があります。

機能リクエスト: ハードウェアがこのオペレーションをサポートしていると思われる場合、または互換モードでもコンパイラに有効なエミュレーション パスがない場合は、機能リクエストを送信してください。通常、オペレーションの前方互換性は保証されています。したがって、カーネルが TPU 世代で実行される場合、将来のすべての世代で実行されますが、以前の世代のエミュレーションが保証されるわけではありません（一部の世代ではキャストのコストが非常に高くなります）。