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ı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

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ığı

görüşmeler

jit_compile=True kullanarak TF'den XLA kullanma

XLA'ya Genel Bakış