Penyesuaian otomatis persisten (khusus GPU)

Kami menggunakan OpenAI Triton untuk menghasilkan beberapa kernel GPU. Triton memungkinkan pembuatan kernel GPU cepat untuk fusion tertentu, tetapi kita harus menyesuaikan beberapa parameter untuk setiap fusi tersebut.

Proses ini dapat memerlukan waktu lama jika ada banyak penggabungan, jadi kami menyediakan cara untuk memuat hasil penyesuaian otomatis tersebut, sambil tetap menjalankan langkah-langkah kompilasi lainnya seperti biasa. Cache penyesuaian otomatis masih berguna jika kita membuat beberapa perubahan: penggabungan yang ada dalam cache akan menggunakan cache, dan cache lainnya akan otomatis disesuaikan secara normal.

Hasil penyesuaian otomatis dapat dibuang/dimuat menggunakan parameter berikut:

--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=

Jika kami menentukan file .txt atau .textproto, cache akan dibuang dalam format textproto, atau dalam format protobuf biner.

Dalam pengujian

Penyesuaian otomatis persisten juga dapat digunakan dalam pengujian. Sebaiknya gunakan metode ini jika pengujiannya sangat besar, terutama jika performa lingkungan pengujian terbatas.

Ini hanya berfungsi dengan baik jika cache autotune berisi hasil yang dihasilkan pada jenis GPU yang sama tempat pengujian dijalankan.

Melakukan penggunaan pengujian dengan penyempurnaan otomatis yang dipertahankan

Untuk saat ini, anggaplah pengujian yang dimaksud selalu menggunakan jenis GPU yang sama.

  1. Kita harus mengekspor hasil autotune dari pengujian, misalnya dengan menentukan parameter ini ke perintah pengujian:

    --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto
    --test_sharding_strategy=disabled
    

    Sharding harus dinonaktifkan untuk mendapatkan satu cache autotune dengan benar untuk semua pengujian.

  2. Kemudian, kita harus mengupload cache tersebut ke repositori kode.

  3. Kemudian, kita harus menambahkan cache ke dependensi data target pengujian, dan memuatnya menggunakan variabel lingkungan.

    data = ["test_autotune_cache.textproto"],
    env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" +
                        "$(execpath test_autotune_cache.textproto)"},
    

    (Anda dapat menggunakan sharding dalam pengujian yang memuat hasil penyesuaian otomatis.)

Lihat juga contoh pengujian di 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

Keusangan cache

Jika banyak perubahan dilakukan pada model, ada kemungkinan cache tidak lagi berisi semua fusion, sehingga pengujian akan menjadi lebih lambat. Dalam hal ini, kita harus membuat ulang cache autotuning.

Jika kita mulai menggunakan jenis GPU baru untuk menjalankan pengujian, hal yang sama berlaku.

Cache juga dapat menjadi tidak terpakai lagi jika compiler XLA berkembang dan menghasilkan fusi yang berbeda.