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

OpenXLA, potansiyel olarak hiçbir kaynak kodu değişikliği olmadan TensorFlow modellerini hızlandırabilen, doğrusal cebir için alana özgü bir derleyicidir.

Giriş

Bir TensorFlow programı çalıştırıldığında, tüm işlemler TensorFlow yürütücüsü tarafından tek tek yürütülür. Her TensorFlow işleminde, yürütücünün dağıtım yaptığı, önceden derlenmiş bir GPU çekirdek uygulaması vardır.

XLA, alternatif bir model çalıştırma modu sağlar: TensorFlow grafiğini, belirtilen model için özel olarak oluşturulan bir işlem çekirdekleri dizisinde 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ığı optimizasyona göz atalı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, biri toplama ve diğeri azaltma için. Ancak XLA, sonucu tek bir çekirdek başlatmada hesaplayacak şekilde grafiği optimize edebilir. Bunu, toplama, çarpma ve azaltma işlemlerini tek bir GPU çekirdeğinde "birleştirerek" gerçekleştirir. Dahası, bu çok kaynaklı 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ılara "aktarır" ve tüm bunları GPU kayıtlarında tutar. Füzyon, XLA'nın en önemli optimizasyonudur. Bellek bant genişliği genellikle donanım hızlandırıcılarda en dar kaynaktır. Bu nedenle, bellek işlemlerinin kaldırılması performansı artırmanın en iyi yollarından biridir.

TensorFlow modelleri için XLA'yı etkinleştir

tf.function(jit_compile=True) ile açık derleme

Açık derleme API'si, hangi işlevlerin derleneceğini seçmek için ayrıntılı bir kontrol sunar. Örneğin, MNIST eğitimini gerçekleştiren aşağıdaki TensorFlow işlevi XLA ile derlenir:

@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'nin derlenmesi gerekir semantiği: Tüm işlevin XLA ile derlenmesi veya errors.InvalidArgumentError istisnası atılır. XLA, boyutların çıkarılabilir olmadığı işlevleri şu anda derleyemez. Diğer bir deyişle, hesaplamanın tamamını çalıştırmadan tüm tensörlerin boyutlarının tahmin edilmesi mümkün değildir. Örneğin, aşağıdaki işlev derlemez:

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

Bununla birlikte, şekiller çeşitli çalıştırmalar arasında farklılık gösterebilir:

@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 ortak çalışmasına ve jit_compile=True kullanımıyla ilgili eğitim videosuna bakın.

Keras ile kullanım

Keras modellerinde jit_compile=True, model.compile öğesine bağımsız değişken olarak ayarlanabilir:

model.compile(optimizer="adam", jit_compile=True)

Dağıtılmış stratejiyle kullanım

XLA:GPU, jit_compile=True ile adım işlevine ek açıklama ekleyerek TF dağıtılmış stratejiyle (MirroredStrategy veya MultiWorkerMirroredStrategy) birlikte 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, TensorFlow işlevleri içinde XLA kullanılarak derlenip yürütülebilen 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 --tf_xla_cpu_global_jit işaretini kullanarak da CPU üzerinde etkinleştirilebilir:

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

Ayrıntılı kullanım örneği için otomatik kümeleme eğiticisi colab'ine bakın.

tfcompile ile CPU için AOT (önceden) derlemesi

Ayrıca, TensorFlow grafiğini yürütülebilir koda dönüştüren (yalnızca x86-64 CPU için) bağımsız bir tfcompile aracı da kullanabilirsiniz.

Derlenmiş programları inceleme

XLA, oluşturulan programları incelemenize olanak tanıyan iç gözlem olanakları sunar. Oluşturulan programların dökümünü almak 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 işlemi gerçekleştirildikten sonra /tmp/generated içinde aşağıdaki dosyaları bulabilirsiniz:

  • module_XXXX.*_optimizations.txt Her derlenen küme başına bir tane olacak şekilde XLA programları oluşturuldu. XLA hata raporlarını gönderirken bu hataları eklemek son derece yararlıdır!

  • module_XXXX.ir-*.ll Dosyalar, NVPTX kaynaklı LLVM orta düzey temsilinde oluşturulur.

  • module_XXXX.ptx PTX dosyaları oluşturuldu.

Ayrıca, XLA kümelerinin TensorFlow grafiğinin içine yerleştirilmesini görselleştiren grafiği şunlarla da oluşturabilirsiniz:

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

Yeniden oluşturulabilir hata raporları

Bir hata raporunun yeniden oluşturulması, oluşturulan XLA programları için dökümler ve kullanılan otomatik kümeleme yerleştirmesi içeriyorsa çok daha kolay olur. Bunları, otomatik kümeleme ile çalışan bir TensorFlow programı için oluşturmak üzere şu komutu çalıştırı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"

Hata bildiriminde bulunurken /tmp/generated dizininin içeriğini ekleyin (yukarıda referans verilmiştir).

Mümkünse run_hlo_module kullanarak ve oluşturulan programlarda tekrar tekrar çalıştırarak bir hatayı tek bir XLA programında tespit etmeye çalışın.

Daha fazla bilgi

XLA Ön Uçları

XLA programları, TensorFlow'dan ayrı olarak aşağıdakiler tarafından oluşturulabilir:

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

Sohbetler

jit_compile=True ile TF'den XLA kullanılıyor

XLA'ya Genel Bakış