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

XLA (Hızlandırılmış Doğrusal Cebir), TensorFlow modellerini potansiyel olarak hiçbir kaynak kodu değişikliği olmadan 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 sunumunda ~7 kat performans artışı ve ~5 kat toplu iş boyutu iyileştirmesi elde edildi:

giriiş

Bir TensorFlow programı çalıştırıldığında, tüm işlemler TensorFlow uygulayı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, modelleri çalıştırmak için alternatif bir mod sağlar: TensorFlow grafiğini, belirli bir model için özel olarak oluşturulmuş bir dizi hesaplama çekirdeği halinde derler. Bu çekirdekler modele özgü olduğundan, optimizasyon için modele özgü bilgileri kullanabilirler. Ö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ışan grafik üç çekirdek başlatır: biri çarpma için, biri toplama için ve biri de azaltma için. Ancak XLA, sonucu tek bir çekirdek başlatmada hesaplayacak şekilde grafiği optimize edebilir. Bunu, toplama, çarpma ve indirgemeyi tek bir GPU çekirdeğinde "kaynaştırarak" yapar. Üstelik bu kaynaştırma işlemi, 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". Fusion, XLA'nın en önemli optimizasyonudur. Bellek bant genişliği genellikle donanım hızlandırıcılardaki 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'sinin derlemesi gereken semantiği vardır: ya tüm işlev XLA ile derlenir ya da bir errors.InvalidArgumentError istisnası atılır. XLA şu anda boyutların tahmin edilemez 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 derlenmez:

@tf.function
def not_compilable(x):
  return tf.unique(x)

Şekiller, çalıştırmalar arasında değişiklik gösterebilir ancak:

@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 öğretici colab'a ve jit_compile=True kullanım hakkında bir eğitim videosuna bakın.

Keras ile kullanım

Keras modelleri için jit_compile=True model.compile için bir bağımsız değişken 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 eklenerek TF dağıtılmış strateji ( MirroredStrategy veya MultiWorkerMirroredStrategy ) ile 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 XLA'yı 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ç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 ayrıca --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 colab'a bakın.

tfcompile ile CPU için AOT (zamanından önce) derlemesi

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ı incelemenizi sağlayan 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

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

  • module_XXXX.*_optimizations.txt Oluşturulan XLA programları , derlenen her küme için bir tane. XLA hata raporları gönderirken bunları eklemek son derece faydalıdır!

  • module_XXXX.ir-*.ll LLVM ara temsilinde, NVPTX gerçekleriyle 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 şu şekilde 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ırması için dökümler içeriyorsa, bir hata raporunun çoğaltılması çok daha kolaydır. Bunları otomatik kümeleme ile çalışan bir TensorFlow programı için 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ı dosyalarken, /tmp/generated dizininin (yukarıda atıfta bulunulan) içeriğini ekleyin.

Mümkünse, run_hlo_module 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 dışında, XLA programları şu şekilde oluşturulabilir:

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

Konuşmalar

jit_compile=True kullanarak TF'den XLA kullanma

XLA'ya Genel Bakış