カテゴリ: コンパイル時間: ホスト オフロード出力の不一致
このエラーは、ホストメモリに明示的にオフロードされたテンソルがプログラム出力として返されたものの、プログラムの出力シグネチャがホストメモリを想定するように構成されていない場合に発生します。
エラー メッセージの例:
INVALID_ARGUMENT: Tensor which is moved to host (starting from tuple.64) is returned from the entry computation but the layout for this output is not set to host memory.
XLA バックエンド: TPU、GPU
概要
コンパイラがテンソルをホスト（CPU）にオフロードするアノテーションを検出すると、次の 3 つのイベントのいずれかが発生するまで、計算グラフを通じてそのテンソルの位置を追跡します。
- デバイスに移動: 一致するアノテーションにより、テンソルがアクセラレータに戻されます。
- ホスト コンピューティング: テンソルがホスト側のオペレーションで使用されます。
- プログラムの終了: テンソルがプログラムの最後に到達し、出力になります。
このエラーは、シナリオ #3 でトリガーされます。テンソルは実行の最後にホストメモリに物理的に配置されますが、XLA プログラムのエントリ計算シグネチャは、その特定の出力をデバイスメモリに配置されるものとして定義します。コンパイラはエントリ計算のインターフェースを暗黙的に変更できないため、エラーが発生します。
デバッグ
このエラーを解決するには、このテンソルをホストの出力にするつもりだったのか、戻る前にデバイスに戻す必要があったのかを判断します。
ホストで返すことを想定: このテンソルをホストメモリで返す（デバイスへの転送を回避する）ことを明示的に指定する場合は、この特定の出力のエントリ計算の出力メモリ空間を ホストメモリに明示的に設定する必要があります。
デバイスで返す予定: テンソルがデバイスに残るか、プログラムの終了前にデバイスに戻る予定だった場合は、アノテーションが欠落している可能性があります。一致するアノテーションを挿入して、テンソルをデバイスに戻します。
オフロードされたテンソルのソースが不明な場合や、「デバイスに移動」アノテーションが欠落している場所が見つからない場合は、XLA ロギングを使用して命令をトレースします。
- ロギングを有効にする: Google Cloud TPU を使用している場合は、次のフラグ
--vmodule=host_offloader=1を使用してプログラムを再実行します。 - ログを分析する: ログで「トレース」出力を探します。これにより、オフロード命令から始まるテンソルのパスが表示されます。これを使用して、テンソルがデバイスに戻されることなくプログラム境界に達する場所を正確に特定します。