OpenXLA adalah compiler khusus domain untuk aljabar linear yang dapat mempercepat model TensorFlow tanpa perubahan kode sumber.
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 dikirimkan oleh eksekutor.
XLA menyediakan mode alternatif untuk model yang sedang berjalan: mode ini mengompilasi grafik TensorFlow menjadi urutan kernel komputasi yang dihasilkan khusus untuk model tertentu. Karena kernel ini unik untuk model, kernel tersebut 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)
Jika dijalankan tanpa XLA, grafik akan meluncurkan tiga kernel: satu untuk perkalian,
satu untuk penambahan, dan satu untuk pengurangan. Namun, XLA dapat mengoptimalkan grafik agar dapat menghitung hasilnya dalam satu peluncuran kernel. Caranya adalah 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; tetapi "mengalirkan" hasil
komputasi perantara ini langsung ke pengguna mereka sekaligus menyimpannya
sepenuhnya dalam register GPU. Fusion adalah satu-satunya pengoptimalan terpenting dari XLA.
Bandwidth memori biasanya merupakan resource paling langka pada akselerator hardware, sehingga menghapus operasi memori adalah salah satu cara terbaik untuk meningkatkan performa.
Mengaktifkan model XLA for TensorFlow
Kompilasi eksplisit dengan tf.function(jit_compile=True)
API kompilasi eksplisit menawarkan kontrol terperinci untuk memilih fungsi yang akan 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 yang dimensinya tidak
dapat disimpulkan: artinya, jika tidak mungkin 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, bentuk dapat bervariasi di seluruh operasi:
@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 mengetahui contoh penggunaan yang lebih detail, 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 perlu melakukan perubahan 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
:
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
Untuk mengetahui contoh penggunaan yang lebih terperinci, 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 untuk 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 dumping dilakukan, Anda dapat menemukan file berikut di
/tmp/generated
:
module_XXXX.*_optimizations.txt
Program XML 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 menengah LLVM, dengan intrinsik NVPTX.module_XXXX.ptx
File PTX yang dihasilkan.
Anda juga dapat membuang grafik yang memvisualisasikan embedding cluster 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 direproduksi jika menyertakan dump untuk program XLA yang dihasilkan dan embedding pengelompokan otomatis yang digunakan. Guna 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
(yang dirujuk di atas).
Jika memungkinkan, coba isolasi
bug ke satu program XLA menggunakan
run_hlo_module
dan jalankan secara berulang pada program yang dihasilkan.
Bacaan lebih lanjut
- Dokumentasi OpenXLA Dokumentasi OpenXLA
- Masalah Umum Daftar masalah umum pada XLA+TF
- XLA - TensorFlow, Dikompilasi: Baca di Blog Google Developers
- Lihat sumber XDA di GitHub.
Frontend XLA
Selain TensorFlow, program XLA dapat dihasilkan oleh:
- JAX: Transformasi composable program Python+NumPy
- Julia: Bahasa Julia untuk komputasi ilmiah
- PyTorch: Framework PyTorch
- Nx: Library komputasi numerik untuk bahasa pemrograman Elixir