Bu sayfa, Cloud Translation API ile çevrilmiştir.
Switch to English

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: örn. XLA kullanılarak 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şturulmuş 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. Ancak XLA, sonucu tek bir çekirdek başlatmasıyla hesaplayacak şekilde grafiği optimize edebilir. Bunu, toplama, çarpma ve indirgeme işlemlerini 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ı doğrudan kullanıcılarına "aktarırken" bunları tamamen GPU kayıtlarında tutuyor. 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, çalışmalar 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 olmadan 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 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 makalesine bakın .

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

Ayrıca, 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 NVPTX içselleri 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 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 bir hatayı tek bir XLA programında izole etmeye çalışın.

daha fazla okuma

XLA Ön Uçları

TensorFlow'dan ayrı olarak, 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