XLA: Mengoptimalkan Kompilator untuk Pembelajaran Mesin

XLA (Accelerated Linear Algebra) adalah kompiler khusus domain untuk aljabar linier yang dapat mempercepat model TensorFlow tanpa kemungkinan perubahan kode sumber.

Hasilnya perbaikan dalam kecepatan dan penggunaan memori: misalnya dalam Bert MLPerf pengajuan menggunakan 8 Volta V100 GPU menggunakan XLA telah mencapai ~ 7x peningkatan kinerja dan ~ 5x bets peningkatan ukuran:

pengantar

Saat program TensorFlow dijalankan, semua operasi dijalankan satu per satu oleh eksekutor TensorFlow. Setiap operasi TensorFlow memiliki implementasi kernel GPU yang telah dikompilasi sebelumnya yang dikirim oleh eksekutor.

XLA menyediakan mode alternatif untuk menjalankan model: XLA mengompilasi grafik TensorFlow ke dalam urutan kernel komputasi yang dibuat khusus untuk model yang diberikan. Karena kernel ini unik untuk model, mereka dapat mengeksploitasi 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)

Jalankan tanpa XLA, grafik meluncurkan tiga kernel: satu untuk perkalian, satu untuk penambahan dan satu untuk pengurangan. Namun, XLA dapat mengoptimalkan grafik sehingga menghitung hasilnya dalam peluncuran kernel tunggal. Ini dilakukan dengan "menyatukan" penambahan, perkalian, dan pengurangan menjadi satu kernel GPU. Selain itu, operasi menyatu ini tidak menuliskan nilai menengah yang diproduksi oleh y*z dan x+y*z ke memori; alih-alih itu "mengalirkan" hasil perhitungan perantara ini langsung ke pengguna mereka sambil menyimpannya sepenuhnya di register GPU. Fusion adalah satu-satunya pengoptimalan terpenting XLA. Bandwidth memori biasanya merupakan sumber daya paling langka pada akselerator perangkat keras, jadi menghapus operasi memori adalah salah satu cara terbaik untuk meningkatkan kinerja.

Aktifkan XLA untuk model TensorFlow

Kompilasi eksplisit dengan tf.function(jit_compile=True)

API kompilasi eksplisit menawarkan kontrol berbutir halus untuk memilih fungsi mana yang harus dikompilasi. Misalnya, fungsi TensorFlow berikut yang menjalankan 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))

The jit_compile API memiliki semantik harus-kompilasi: baik seluruh fungsi dikompilasi dengan XLA, atau errors.InvalidArgumentError eksepsi dilemparkan. XLA saat ini tidak dapat mengkompilasi fungsi di mana dimensi tidak inferrable: yaitu, jika itu tidak mungkin untuk menyimpulkan dimensi semua tensor tanpa menjalankan seluruh perhitungan. Misalnya, fungsi berikut tidak akan dikompilasi:

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

Bentuk dapat bervariasi di seluruh proses:

@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 penggunaan contoh yang lebih rinci, dan tutorial video tentang jit_compile=True penggunaan.

Pengelompokan otomatis

Cara mudah untuk mulai menggunakan XLA dalam model TensorFlow tanpa perubahan apapun untuk mengaktifkan auto-clustering, yang secara otomatis menemukan cluster (subgraphs terhubung) dalam fungsi TensorFlow yang dapat dikompilasi dan dieksekusi menggunakan XLA. Auto-clustering GPU dapat diaktifkan dengan menetapkan TF_XLA_FLAGS variabel lingkungan:

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

Auto-pengelompokan saat ini dioptimalkan untuk beban kerja GPU, tetapi juga dapat diaktifkan pada CPU dengan tambahan menggunakan bendera --tf_xla_cpu_global_jit :

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

Untuk contoh penggunaan rinci melihat auto-pengelompokan tutorial colab .

AOT (Menjelang-of-time) kompilasi untuk CPU dengan tfcompile

Anda juga dapat menggunakan mandiri tfcompile alat, yang mengubah TensorFlow grafik ke dalam kode executable (untuk CPU x86-64 saja).

Periksa program yang dikompilasi

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

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

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

  • module_XXXX.*_optimizations.txt Generated program XLA , satu per setiap cluster disusun. Melampirkan itu saat mengirimkan laporan bug XLA sangat membantu!

  • module_XXXX.ir-*.ll file Generated di LLVM representasi menengah, dengan NVPTX intrinsik.

  • module_XXXX.ptx Generated PTX file.

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

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

Laporan bug yang dapat direproduksi

Laporan bug jauh lebih mudah untuk 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 mengajukan bug, melampirkan isi dari /tmp/generated direktori (dirujuk di atas).

Jika memungkinkan, cobalah untuk mengisolasi bug untuk program XLA tunggal dengan menggunakan replay_computation dan iteratif menjalankannya pada program yang dihasilkan.

Bacaan lebih lanjut

Tampilan Depan XLA

Selain TensorFlow, program XLA dapat dibuat dengan:

  • JAX : Composable transformasi dari program Python + NumPy
  • Julia : Bahasa Julia untuk komputasi ilmiah
  • PyTorch kerangka PyTorch:
  • Nx : Numerical komputasi perpustakaan untuk bahasa pemrograman Elixir

Pembicaraan

Menggunakan XLA dari TF menggunakan jit_compile=True

Ikhtisar XLA