XLA (Hızlandırılmış Doğrusal Cebir), potansiyel olarak kaynak kodu değişikliği olmadan TensorFlow modellerini hızlandırabilen doğrusal cebir için alana özgü bir derleyicidir.
Sonuçlar hız ve bellek kullanımında iyileştirmelerdir: örneğin, XLA kullanan 8 Volta V100 GPU kullanan BERT MLPerf gönderiminde ~7x performans iyileştirmesi ve ~5x parti boyutu iyileştirmesi elde edilmiştir:

giriiş
Bir TensorFlow programı çalıştırıldığında, tüm işlemler TensorFlow yürütücüsü tarafından ayrı ayrı yürütülür. Her TensorFlow işlemi, yürütücünün gönderdiği önceden derlenmiş bir GPU çekirdek uygulamasına sahiptir.
XLA, çalışan modeller için alternatif bir mod sağlar: TensorFlow grafiğini, belirli bir model için özel olarak oluşturulmuş bir dizi hesaplama çekirdeğinde derler. Bu çekirdekler modele özgü olduğundan, optimizasyon için modele özgü bilgilerden yararlanabilirler. Örneğin, basit bir TensorFlow hesaplaması bağlamında XLA'nın yaptığı optimizasyona bakalım:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
XLA olmadan çalıştırıldığında, grafik üç çekirdeği başlatır: biri çarpma için, biri toplama için ve biri de indirgeme için. Ancak XLA, sonucu tek bir çekirdek başlatmasında hesaplayacak şekilde grafiği optimize edebilir. Bunu, toplama, çarpma ve küçültmeyi tek bir GPU çekirdeğinde "birleştirerek" yapar. Ayrıca, bu birleştirilmiş işlem, y*z
ve x+y*z
tarafından üretilen ara değerleri belleğe yazmaz; bunun yerine, bu ara hesaplamaların sonuçlarını, tamamen GPU kayıtlarında tutarken, doğrudan kullanıcılarına "aktarır". Fusion, XLA'nın en önemli optimizasyonudur. Bellek bant genişliği genellikle donanım hızlandırıcılarda en kıt kaynaktır, bu nedenle bellek işlemlerini kaldırmak performansı artırmanın en iyi yollarından biridir.
TensorFlow modelleri için XLA'yı etkinleştirin
tf.function(jit_compile=True)
ile açık derleme
Açık derleme API'si, hangi işlevlerin derlenmesi gerektiğini seçmek için ayrıntılı bir denetim sunar. Örneğin, MNIST eğitimini gerçekleştiren aşağıdaki TensorFlow işlevi XLA ile derlenmiştir:
@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, zorunlu derleme semantiğine sahiptir: ya tüm işlev XLA ile derlenir ya da bir errors.InvalidArgumentError
istisnası atılır. XLA şu anda boyutların çıkarsanamadığı işlevleri derleyemez : yani, tüm hesaplamayı çalıştırmadan tüm tensörlerin boyutlarını çıkarmak mümkün değilse. Örneğin, aşağıdaki işlev derlenmeyecektir:
@tf.function
def not_compilable(x):
return tf.unique(x)
Şekiller, koşular arasında değişebilir:
@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]))
Daha ayrıntılı bir kullanım örneği için eğitim işbirliğine ve jit_compile=True
ile ilgili bir eğitim videosuna bakın.
Keras ile kullanım
Keras modelleri için jit_compile=True
, model.compile
için bir argüman olarak ayarlanabilir:
model.compile(optimizer="adam", jit_compile=True)
Dağıtılmış strateji ile kullanım
XLA:GPU, jit_compile=True
ile adım işlevine açıklama ekleyerek TF dağıtılmış stratejisiyle ( MirroredStrategy
veya MultiWorkerMirroredStrategy
) kullanılabilir:
@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)
Otomatik kümeleme
TensorFlow modellerinde herhangi bir değişiklik yapmadan XLA'yı kullanmaya başlamanın basit bir yolu, XLA kullanılarak derlenip yürütülebilen TensorFlow işlevleri içinde kümeleri (bağlı alt grafikler) otomatik olarak bulan otomatik kümelemeyi etkinleştirmektir. GPU'da otomatik kümeleme, TF_XLA_FLAGS
ortam değişkeni ayarlanarak etkinleştirilebilir:
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
Otomatik kümeleme şu anda GPU iş yükleri için optimize edilmiştir, ancak ek olarak --tf_xla_cpu_global_jit
bayrağı kullanılarak CPU'da da etkinleştirilebilir:
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
Ayrıntılı bir kullanım örneği için otomatik kümeleme öğretici işbirliğine bakın .
tfcompile
ile CPU için AOT (Zamanından önce) derlemesi
Ayrıca TensorFlow grafiğini yürütülebilir koda dönüştüren bağımsız bir tfcompile
aracı da kullanabilirsiniz (yalnızca x86-64 CPU için).
Derlenmiş programları inceleyin
XLA, oluşturulan programları incelemenize izin veren iç gözlem olanakları sağlar. Oluşturulan programları boşaltmak için XLA_FLAGS
ortam değişkenini kullanın:
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
Damping gerçekleştirildikten sonra, aşağıdaki dosyaları /tmp/generated
içinde bulabilirsiniz:
module_XXXX.*_optimizations.txt
Derlenen her küme için bir tane olmak üzere oluşturulan XLA programları . XLA hata raporlarını gönderirken bunları eklemek son derece yararlıdır!module_XXXX.ir-*.ll
ile LLVM ara gösteriminde oluşturulan dosyalar.module_XXXX.ptx
Oluşturulan PTX dosyaları.
Ayrıca, XLA kümelerinin TensorFlow grafiğinin içine gömülmesini görselleştiren grafiği aşağıdakilerle de dökebilirsiniz:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Tekrarlanabilir hata raporları
Oluşturulan XLA programları ve kullanılan otomatik kümeleme katıştırma için dökümler içeriyorsa, bir hata raporunun çoğaltılması çok daha kolaydır. Otomatik kümeleme ile çalışan bir TensorFlow programı için bunları oluşturmak için şunu başlatın:
$ 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"
Hataları dosyalarken, /tmp/generated
dizininin içeriğini ekleyin (yukarıda referans alınmıştır).
Mümkünse, replay_computation
kullanarak ve onu oluşturulan programlarda yinelemeli olarak çalıştırarak bir hatayı tek bir XLA programından ayırmaya çalışın.
daha fazla okuma
- Bilinen Sorunlar XLA ile ilgili bilinen sorunların listesi
- XLA Mimarisi : XLA mimarisine genel bakış
- XLA - TensorFlow, Derlenmiş : Google Developers Blog'da Okuyun
- Github'daki XLA kaynağına göz atın!
XLA Ön Uçları
TensorFlow dışında, XLA programları aşağıdakiler tarafından oluşturulabilir:
- JAX : Python+NumPy programlarının birleştirilebilir dönüşümleri
- Julia : Bilimsel hesaplama için Julia dili
- PyTorch : PyTorch çerçevesi
- Nx : Elixir programlama dili için sayısal hesaplama kitaplığı