Bir sorunuz mu var? TensorFlow Forum Ziyaret Forumunda toplulukla bağlantı kurun

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.

Sonuçlar hız ve bellek kullanımındaki iyileştirmelerdir: örneğin XLA kullanarak 8 Volta V100 GPU kullanan BERT MLPerf gönderiminde ~ 7 kat performans artışı ve ~ 5 kat toplu iş boyutu iyileştirmesi elde edildi:

Giriş

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 model için özel olarak oluşturulan bir dizi hesaplama çekirdeğinde derler. Bu çekirdekler modele özel olduğundan, optimizasyon için modele özgü bilgilerden yararlanabilirler. Örneğin, XLA'nın basit bir TensorFlow hesaplaması bağlamında yaptığı bir optimizasyona bakalım:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

XLA olmadan çalıştırın, grafik üç çekirdek başlatır: biri çarpma, biri toplama ve diğeri azaltma için. Bununla birlikte, XLA grafiği, sonucu tek bir çekirdek başlatmasıyla hesaplayacak şekilde optimize edebilir. Bunu, toplama, çarpma ve indirgeme işlemlerini tek bir GPU çekirdeğine "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ı doğrudan kullanıcılarına "aktarır" ve bunları tamamen GPU kayıtlarında tutar. Fusion, XLA'nın en önemli optimizasyonudur. Bellek bant genişliği tipik olarak donanım hızlandırıcılarda en az bulunan 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 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'sinin derlenmesi gereken semantiği vardır: ya işlevin tamamı XLA ile derlenir ya da bir errors.InvalidArgumentError istisnası atılır. XLA şu anda boyutların çıkarılamaz olduğu 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ğitici makaleye bakın.

Otomatik kümeleme

XLA'yı TensorFlow modellerinde herhangi bir değişiklik yapmadan kullanmaya başlamanın basit bir yolu, XLA kullanılarak derlenebilen ve yürütülebilen TensorFlow işlevleri içindeki 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 CPU'da ek olarak --tf_xla_cpu_global_jit bayrağı kullanılarak da --tf_xla_cpu_global_jit :

$ 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 makalesine bakın .

tfcompile ile CPU için AOT (Ahead-of-time) derlemesi

TensorFlow grafiğini çalıştırılabilir 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ı incelemenizi sağlayan iç gözlem olanakları sağlar. Oluşturulan programları dökmek için, XLA_FLAGS ortam değişkenini XLA_FLAGS :

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

Döküm gerçekleştirildikten sonra, aşağıdaki dosyaları /tmp/generated içinde bulabilirsiniz:

  • module_XXXX.*_optimizations.txt Oluşturulan XLA programları , her derlenen küme için bir tane. 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, XLA kümelerinin TensorFlow grafiğinin içine gömülmesini görselleştiren grafiği aşağıdakilerle boşaltabilirsiniz:

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

Tekrarlanabilir hata raporları

Oluşturulan XLA programları için dökümler ve kullanılan otomatik kümeleme yerleştirme içeriyorsa bir hata raporunun yeniden oluşturulması çok daha kolaydır. Otomatik kümeleme ile çalışan bir TensorFlow programı için bunları oluşturmak üzere ş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ı doldururken, /tmp/generated dizinin içeriğini ekleyin (yukarıda atıfta bulunulmuştur).

Mümkünse, replay_computation kullanarak ve oluşturulan programlarda yinelemeli olarak çalıştırarak tek bir XLA programında bir hatayı izole etmeye çalışın.

daha fazla okuma

XLA Ön Uçları

TensorFlow dışında, XLA programları şu şekilde oluşturulabilir:

  • JAX : Python + NumPy programlarının bir araya getirilebilir dönüşümleri
  • Julia : Bilimsel bilgi işlem için Julia dili
  • PyTorch : PyTorch çerçevesi
  • Nx : Elixir programlama dili için sayısal hesaplama kitaplığı