XLA: Mengoptimalkan Compiler untuk Machine Learning

OpenXLA adalah compiler khusus domain untuk aljabar linear yang dapat mempercepat model TensorFlow tanpa kemungkinan perubahan kode sumber.

Pengantar

Saat program TensorFlow dijalankan, semua operasi dijalankan secara terpisah oleh eksekutor TensorFlow. Setiap operasi TensorFlow memiliki implementasi kernel GPU yang telah dikompilasi sebelumnya yang dikirimkan oleh eksekutor.

XLA menyediakan mode alternatif untuk menjalankan model: XLA mengompilasi grafik TensorFlow menjadi urutan kernel komputasi yang dihasilkan secara khusus untuk model tertentu. Karena kernel ini unik untuk model, kernel ini dapat memanfaatkan informasi khusus model untuk pengoptimalan. Misalnya, mari kita lihat pengoptimalan yang dilakukan XLA dalam konteks komputasi TensorFlow sederhana:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

Jika dijalankan tanpa XLA, grafik akan meluncurkan tiga kernel: satu untuk perkalian, satu untuk penambahan, dan satu untuk pengurangan. Namun, XLA dapat mengoptimalkan grafik sehingga menghitung hasilnya dalam satu peluncuran kernel. Hal ini dilakukan dengan "menggabungkan" penambahan, perkalian, dan pengurangan menjadi satu kernel GPU. Selain itu, operasi gabungan ini tidak menulis nilai perantara yang dihasilkan oleh y*z dan x+y*z ke memori; sebagai gantinya, operasi ini "menstreaming" hasil komputasi perantara ini langsung ke penggunanya sekaligus menyimpannya sepenuhnya di register GPU. Fusi adalah pengoptimalan terpenting XLA. Bandwidth memori biasanya merupakan resource yang paling langka di akselerator hardware, sehingga menghapus operasi memori adalah salah satu cara terbaik untuk meningkatkan performa.

Mengaktifkan XLA untuk model TensorFlow

Kompilasi eksplisit dengan tf.function(jit_compile=True)

API kompilasi eksplisit menawarkan kontrol terperinci untuk memilih fungsi yang harus dikompilasi. Misalnya, fungsi TensorFlow berikut yang melakukan pelatihan MNIST dikompilasi dengan XLA:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

jit_compile API memiliki semantik harus dikompilasi: seluruh fungsi dikompilasi dengan XLA, atau pengecualian errors.InvalidArgumentError ditampilkan. XLA saat ini tidak dapat mengompilasi fungsi jika dimensinya tidak dapat disimpulkan: yaitu, jika tidak dapat menyimpulkan dimensi semua tensor tanpa menjalankan seluruh komputasi. Misalnya, fungsi berikut tidak akan dikompilasi:

@tf.function
def not_compilable(x):
  return tf.unique(x)

Namun, bentuknya dapat bervariasi di seluruh run:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

Lihat colab tutorial untuk contoh penggunaan yang lebih mendetail, dan video tutorial tentang penggunaan jit_compile=True.

Penggunaan dengan Keras

Untuk model Keras, jit_compile=True dapat ditetapkan sebagai argumen ke model.compile:

model.compile(optimizer="adam", jit_compile=True)

Penggunaan dengan strategi terdistribusi

XLA:GPU dapat digunakan dengan strategi terdistribusi TF (MirroredStrategy atau MultiWorkerMirroredStrategy) dengan menganotasi fungsi langkah dengan jit_compile=True:

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

Pengelompokan otomatis

Cara mudah untuk mulai menggunakan XLA dalam model TensorFlow tanpa perubahan apa pun adalah dengan mengaktifkan pengelompokan otomatis, yang secara otomatis menemukan cluster (subgrafik yang terhubung) dalam fungsi TensorFlow yang dapat dikompilasi dan dieksekusi menggunakan XLA. Pengelompokan otomatis di GPU dapat diaktifkan dengan menetapkan variabel lingkungan TF_XLA_FLAGS:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Pengelompokan otomatis saat ini dioptimalkan untuk workload GPU, tetapi juga dapat diaktifkan di CPU dengan menggunakan flag --tf_xla_cpu_global_jit secara tambahan:

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Untuk contoh penggunaan mendetail, lihat colab tutorial pengelompokan otomatis.

Kompilasi AOT (Ahead-of-time) untuk CPU dengan tfcompile

Anda juga dapat menggunakan alat tfcompile mandiri, yang mengonversi grafik TensorFlow menjadi kode yang dapat dieksekusi (khusus CPU x86-64).

Memeriksa program yang dikompilasi

XLA menyediakan fasilitas introspeksi yang memungkinkan Anda memeriksa program yang dihasilkan. Untuk membuang program yang dihasilkan, gunakan variabel lingkungan XLA_FLAGS:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

Setelah pembuangan dilakukan, Anda dapat menemukan file berikut di /tmp/generated:

  • module_XXXX.*_optimizations.txt Program XLA yang dihasilkan, satu per setiap cluster yang dikompilasi. Melampirkannya saat mengirimkan laporan bug XLA sangat membantu.

  • module_XXXX.ir-*.ll File yang dihasilkan dalam representasi perantara LLVM, dengan intrinsik NVPTX.

  • module_XXXX.ptx File PTX yang dihasilkan.

Anda juga dapat membuang grafik yang memvisualisasikan penyematan cluster XLA di dalam grafik TensorFlow dengan:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

Laporan bug yang dapat direkonstruksi

Laporan bug jauh lebih mudah direproduksi jika menyertakan dump untuk program XLA yang dihasilkan dan penyematan pengelompokan otomatis yang digunakan. Untuk membuatnya untuk program TensorFlow yang berjalan dengan pengelompokan otomatis, luncurkan:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

Saat melaporkan bug, lampirkan konten direktori /tmp/generated (dirujuk di atas).

Jika memungkinkan, coba isolasi bug ke satu program XLA menggunakan run_hlo_module dan jalankan secara iteratif pada program yang dihasilkan.

Bacaan lebih lanjut

Frontend XLA

Selain TensorFlow, program XLA dapat dibuat oleh:

  • JAX: Transformasi composable dari program Python+NumPy
  • Julia: Bahasa Julia untuk komputasi ilmiah
  • PyTorch: Framework PyTorch
  • Nx: Library komputasi numerik untuk bahasa pemrograman Elixir

Bincang-bincang

Menggunakan XLA dari TF menggunakan jit_compile=True

Ringkasan XLA