GPU カーネルの生成には OpenAI Triton が使用されています。Triton を使用すると、特定の融合に対して高速 GPU カーネルを生成できますが、そのような融合ごとに一部のパラメータを調整する必要があります。
融合が多い場合、この処理には長い時間がかかる可能性があります。そのため、他のコンパイル ステップを通常どおり実行しながら、これらの自動チューニングの結果を読み込む方法を用意しています。自動チューニング キャッシュは、いくつかの変更を加える場合に役立ちます。キャッシュに存在する融合はキャッシュを使用し、他の融合は通常どおりに自動調整されます。
自動チューニングの結果は、次のパラメータを使用してダンプ/読み込みできます。
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
.txt ファイルまたは .textproto ファイルを指定すると、キャッシュは textproto 形式でダンプされます。それ以外の場合は、バイナリ protobuf 形式でダンプされます。
テスト中
テストでは、永続的な自動チューニングも使用できます。テストが非常に大規模な場合、特にテスト環境のパフォーマンスが制限されている場合は、この方法をおすすめします。
自動チューニング キャッシュに、テストを実行しているのと同じタイプの GPU で生成された結果が含まれている場合にのみ適切に機能します。
永続自動チューニングを使用したテスト
ここでは、対象のテストで常に同じ GPU タイプが使用されていると仮定します。
たとえば、次のパラメータをテストコマンドに指定するなどして、テストから自動チューニングの結果をエクスポートする必要があります。
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledすべてのテストに対して 1 つの自動チューニング キャッシュを正しく取得するには、シャーディングを無効にする必要があります。
次に、そのキャッシュをコード リポジトリにアップロードする必要があります。
次に、テスト ターゲットのデータ依存関係にキャッシュを追加し、環境変数を使用してキャッシュを読み込む必要があります。
data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"},（自動チューニングの結果を読み込むテストではシャーディングを使用しても問題ありません）。
xla/service/gpu/tests/BUILD にあるテストの例もご覧ください。
- load_autotune_results_using_execpath_test
- load_autotune_results_from_test_workspace_test
- dump_autotune_results_to_test_outputs_test
キャッシュの廃止
モデルに多数の変更を加えると、キャッシュにすべての融合が含まれなくなるため、テストが遅くなる可能性があります。この場合、自動チューニング キャッシュを再生成する必要があります。
テストの実行に新しいタイプの GPU を使用する場合も、同じことが言えます。
また、XLA コンパイラが進化して異なる融合が生成される場合も、キャッシュが古くなることがあります。