XLA: Makine Öğrenimi için Derleyiciyi Optimize Etme

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.

Örneğin Bert: Sonuçlar hızı ve hafıza kullanımı iyileştirmeler MLPerf bir ~ 7x performans artışı elde ve ~ 5x parti boyutu gelişme vardır XLA kullanılarak 8 Volta V100 GPU'larda kullanılarak sunulması:

Tanıtım

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 çekirdeği 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şturulan 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ığı bir 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 erimiş işlemi ile üretilen ara değerler yazmak etmez y*z ve x+y*z bellek; 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

İle Açık derleme tf.function(jit_compile=True)

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 mutlaka derleme anlambilim içeriyor: ya tüm fonksiyon XLA ile derlendi veya bir errors.InvalidArgumentError istisnası atılır. Olduğunu, tüm hesaplama çalıştırmadan tüm tensörlerinin boyutlarını ortaya koymak mümkün değilse: XLA şu anda boyutlar inferrable olmayan işlevleri derleme olamaz. Ö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]))

Bkz öğretici CoLab daha detaylı kullanım örneği için, ve bir öğretici bir video üzerinde jit_compile=True kullanımını.

Otomatik kümeleme

Basit bir şekilde, otomatik olarak derlenmiş ve XLA kullanılarak yürütülebilir TensorFlow fonksiyonları olan kümeleri (bağlı subgraphs) bulur otomatik kümeleme sağlamak için herhangi bir değişiklik olmadan TensorFlow modellerinde XLA kullanılarak başlatın. GPU üzerindeki otomatik kümeleme ayarlayarak etkinleştirilebilir TF_XLA_FLAGS ortam değişkeni:

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

Ayrıca bayrak kullanarak otomatik kümeleme anda GPU iş yükleri için optimize edilmiş, ancak aynı zamanda CPU üzerinde etkinleştirilebilir --tf_xla_cpu_global_jit :

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

Daha detaylı bir kullanım Örneğin bkz auto-kümeleme öğretici CoLab .

İle CPU için AOT (Öncesinde-of-time) derleme tfcompile

Ayrıca, bir tek başına kullanabilir tfcompile (X86-64 CPU için) çalıştırılabilir kod içine TensorFlow grafik dönüştüren bir araç.

Derlenmiş programları inceleyin

XLA, oluşturulan programları incelemenize izin veren iç gözlem olanakları sağlar. Oluşturulan programlar dökümü için ortam değişkeni kullanmak XLA_FLAGS :

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

Gerçekleştirilir damping sonra, ilerleyen dosyaları bulabilirsiniz /tmp/generated :

  • module_XXXX.*_optimizations.txt Oluşturulan XLA programları , her derlenmiş küme başına bir. XLA hata raporlarını gönderirken bunları eklemek son derece yararlıdır!

  • module_XXXX.ir-*.ll Oluşan dosyaları LLVM ile, ara temsil NVPTX intrinsics.

  • module_XXXX.ptx Oluşturulan PTX dosyaları.

Ayrıca, TensorFlow grafiğinin içine XLA kümelerinin 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"

Böcek dosyalama yaparken, içeriğini eklemek /tmp/generated (yukarıda başvurulan) dizinine.

Mümkünse, kullanarak tek XLA programına bir hata izole etmeye replay_computation oluşturulan programlarla sitede yayınlayarak iteratif ve.

daha fazla okuma

XLA Ön Uçları

TensorFlow dışında, XLA programları aşağıdakiler tarafından oluşturulabilir:

  • JAX : Python + NumPy programlarının composable dönüşümleri
  • Julia : bilimsel hesaplama için Julia dil
  • PyTorch : PyTorch çerçevesi
  • Nx : Elixir programlama dili için Sayısal hesaplama kütüphanesi

görüşmeler

Kullanarak TF den XLA kullanma jit_compile=True

XLA'ya Genel Bakış