Terima kasih telah mendengarkan Google I/O. Lihat semua sesi sesuai permintaan Tonton sesuai permintaan

XLA: Mengoptimalkan Kompiler untuk Pembelajaran Mesin

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

Hasilnya adalah peningkatan dalam kecepatan dan penggunaan memori: misalnya dalam pengiriman BERT MLPerf menggunakan 8 Volta V100 GPU menggunakan XLA telah mencapai peningkatan kinerja ~7x dan peningkatan ukuran batch ~5x:

Perkenalan

Saat program TensorFlow dijalankan, semua operasi dijalankan secara individual oleh pelaksana TensorFlow. Setiap operasi TensorFlow memiliki implementasi kernel GPU yang telah dikompilasi yang dikirim oleh eksekutor.

XLA menyediakan mode alternatif untuk menjalankan model: XLA mengkompilasi grafik TensorFlow ke dalam urutan kernel komputasi yang dibuat khusus untuk model tertentu. Karena kernel ini unik untuk model, mereka dapat mengeksploitasi informasi khusus model untuk pengoptimalan. Sebagai contoh, 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 penjumlahan, dan satu untuk pengurangan. Namun, XLA dapat mengoptimalkan grafik sehingga menghitung hasilnya dalam peluncuran kernel tunggal. Ini dilakukan dengan "menggabungkan" penambahan, perkalian, dan pengurangan ke dalam satu kernel GPU. Selain itu, operasi gabungan ini tidak menuliskan nilai antara yang dihasilkan oleh y*z dan x+y*z ke memori; alih-alih "mengalirkan" hasil perhitungan menengah ini langsung ke penggunanya sambil menyimpannya sepenuhnya dalam register GPU. Fusion adalah satu-satunya pengoptimalan XLA yang paling penting. 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 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))

API jit_compile memiliki semantik yang harus dikompilasi : seluruh fungsi dikompilasi dengan XLA, atau pengecualian errors.InvalidArgumentError dilempar. XLA saat ini tidak dapat mengkompilasi fungsi di mana dimensi tidak dapat disimpulkan : yaitu, jika tidak memungkinkan untuk menyimpulkan dimensi semua tensor tanpa menjalankan seluruh komputasi. Misalnya, fungsi berikut tidak dapat dikompilasi:

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

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

Penggunaan dengan Keras

Untuk model Keras, jit_compile=True dapat ditetapkan sebagai argumen untuk 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 sederhana untuk mulai menggunakan XLA dalam model TensorFlow tanpa perubahan apa pun adalah mengaktifkan pengelompokan otomatis , yang secara otomatis menemukan kluster (subgraf terhubung) dalam fungsi TensorFlow yang dapat dikompilasi dan dijalankan menggunakan XLA. Pengelompokan otomatis pada GPU dapat diaktifkan dengan menyetel variabel lingkungan TF_XLA_FLAGS :

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

Pengelompokan otomatis saat ini dioptimalkan untuk beban kerja GPU, tetapi juga dapat diaktifkan pada CPU dengan tambahan menggunakan flag --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 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 (hanya untuk CPU x86-64).

Periksa program yang dikompilasi

XLA menyediakan fasilitas introspeksi yang memungkinkan Anda memeriksa program yang dihasilkan. Untuk membuang program yang dibuat, gunakan 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 berikut di /tmp/generated :

  • module_XXXX.*_optimizations.txt Program XLA yang dihasilkan , satu per setiap kluster 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 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, lampirkan konten direktori /tmp/generated (dirujuk di atas).

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

Bacaan lebih lanjut

Frontend XLA

Selain TensorFlow, program XLA dapat dibuat dengan:

  • JAX : Transformasi program Python+NumPy yang dapat disusun
  • Julia : Bahasa Julia untuk komputasi ilmiah
  • PyTorch : kerangka kerja PyTorch
  • Nx : Pustaka komputasi numerik untuk bahasa pemrograman Elixir

Pembicaraan

Menggunakan XLA dari TF menggunakan jit_compile=True

Ikhtisar XLA