Dump HLO Computations

Dump HLO adalah representasi tekstual modul HLO pada berbagai tahap komputasi. Informasi ini berguna untuk proses debug, dan Anda sering kali perlu menyertakannya dalam laporan bug. File ini biasanya berupa file teks yang dapat dibaca manusia dan mencantumkan petunjuk HLO dan propertinya. Terkadang, modul HLO di-dump sebagai:

  • HloProto: File buffer protokol, yang merupakan format yang lebih terstruktur dan dapat dibaca mesin.
  • HloSnapshot: Modul HLO beserta inputnya. Untuk memutar ulang HLO, Anda terkadang memerlukan input aktual yang dimasukkan ke komputasi tertentu, bukan data acak.

Anda dapat menggunakan flag XLA untuk menentukan dan mendapatkan dump. Dalam sebagian besar kasus, Anda dapat menyetelnya dengan variabel lingkungan. JAX juga menawarkan cara terprogram untuk mencetak dump HLO.

Eksekusi Lokal

Menggunakan Variabel Lingkungan

Anda dapat menetapkan variabel lingkungan XLA_FLAGS dengan tanda yang diperlukan untuk mendapatkan dump. Ini berfungsi untuk JAX, TensorFlow, dan PyTorch/XLA.

Untuk mengekspor modul HLO dan informasi debug lainnya ke direktori tertentu, jalankan program Anda dengan tanda --xla_dump_to:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Misalnya, Anda dapat menggunakan /tmp atau /tmp/xladump sebagai jalur.

Secara default, ini akan membuang modul HLO sebagai teks, di awal dan akhir pipeline pengoptimalan.

Anda juga dapat menentukan format secara eksplisit:

  1. Dump teks
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. Proto HLO
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. Snapshot HLO
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Render grafik dengan server graphviz (hanya berfungsi baik untuk grafik kecil)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Render grafik ke file HTML (hanya berfungsi dengan baik untuk grafik kecil)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Untuk grafik yang lebih besar, Anda dapat menggunakan interactive_graphviz untuk memvisualisasikan bagian grafik.

Membuang Kartu Perantara Khusus

Selain HLO yang telah dioptimalkan sebelumnya / dioptimalkan akhir, Anda juga dapat membuang status HLO setelah melewati kompilator tertentu.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

Modul HLO akan di-dump untuk lintasan yang namanya cocok dengan ekspresi reguler (regex). Misalnya, Anda dapat mengamati HLO yang dihasilkan dari lintasan yang terkait dengan partisi SPMD dengan:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

Untuk mengekspor hasil setelah setiap proses XLA (tindakan ini akan menghasilkan banyak file), Anda dapat menetapkan:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

Opsi khusus JAX

Secara terprogram di JAX

Selain meneruskan flag atau variabel lingkungan, Anda juga dapat membuang HLO secara terprogram menggunakan API lower dan compile JAX.

Ambil HLO asli yang tidak dioptimalkan dan diturunkan secara lokal dengan:

jax.jit(f).lower(*args).as_text('hlo')

Untuk melakukan dumping ke file selama proses kompilasi HLO, tentukan:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Dump jaxprs

jaxprs adalah representasi perantara JAX untuk rekaman aktivitas program. Untuk membuat dump ini, tetapkan variabel lingkungan:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Pelajari lebih lanjut di dokumentasi JAX tentang Mengekspor dan menyerialkan komputasi yang di-stage-out: Debugging.

Google Colab

Variabel lingkungan

Di sel pertama yang dieksekusi di notebook Anda (karena variabel lingkungan dan flag command line biasanya hanya diproses satu kali, misalnya, pada waktu impor modul atau waktu inisialisasi backend XLA), tambahkan XLA_FLAGS yang dijelaskan di atas dengan os.environ, misalnya:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

Tindakan ini akan mengalihkan komputasi ke DIRECTORY_PATH, misalnya /tmp. Di Colab, buka browser "Files" di sidebar kiri, untuk melihat dan mengakses direktori ini.

Anda dapat menggunakan semua tanda yang disebutkan di bagian Eksekusi Lokal.

Opsi khusus JAX

Mirip dengan eksekusi lokal; untuk introspeksi interaktif langsung, Anda dapat langsung mencetak HLO pra-pengoptimalan komputasi:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

Anda juga dapat langsung mencetak HLO yang dioptimalkan komputasi:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Menghapus Semua/Sebagian Komputasi

Jika Anda ingin melihat semuanya dalam dump, termasuk semua kompilasi kecil, tetapkan variabel lingkungan JAX:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mozaik

Mosaic adalah compiler untuk backend Pallas TPU, dan backend Pallas GPU eksperimental. Untuk melakukan dump komputasi mosaik, tetapkan flag berikut:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Atau, tetapkan argumen inisialisasi TPU sebagai variabel lingkungan:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Lihat dokumentasi JAX tentang Pallas dan Mosaic untuk mempelajari lebih lanjut.

Lebih lanjut tentang Dump HLO

Menemukan komputasi yang tepat

Biasanya, banyak komputasi yang dibuang. File yang di-dump diberi nama secara eksplisit dengan "nama komputasi" JAX, Tensorflow, atau PyTorch/XLA yang dipanggil dalam log, sehingga memudahkan identifikasi file HLO yang relevan. Contoh:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Atau, Anda dapat menggunakan ripgrep untuk mengidentifikasi dengan cepat modul yang berisi simbol atau komputasi tertentu.

Tips: Sertakan 3 file yang di-dump sebelum/sesudah/penugasan buffer yang relevan dalam laporan bug Anda.

Konversi HLO

Alat bernama hlo-opt yang dapat menerjemahkan antara format HLOProto dan teks. Hal ini berguna jika Anda memiliki satu format, tetapi memerlukan format lainnya untuk proses debug.

Pelajari cara menggunakannya: Dokumentasi Alat XLA: hlo-opt.

Putar ulang

Anda dapat menjalankan (memutar ulang) komputasi yang di-dump pada backend XLA tertentu dengan data palsu atau snapshot input. Cara ini adalah cara yang mudah untuk mereproduksi, melakukan iterasi, dan men-debug masalah di XLA.

Perintah berikut menggunakan data palsu. Jika telah menyimpan Snapshot HLO, Anda dapat meneruskannya, dan data dari snapshot akan digunakan. Untuk tetap menggunakan data palsu saat menjalankan snapshot, teruskan tanda --force_fake_data.

Backend CPU:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

Backend GPU:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt