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

OpenXLA, doğrusal cebir için alana özgü bir derleyicidir ve TensorFlow modellerini kaynak kodunda değişiklik yapmadan hızlandırabilir.

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şleminin, yürütücünün dağıttığı önceden derlenmiş bir GPU çekirdeği uygulaması vardır.

XLA, modelleri çalıştırmak için alternatif bir mod sağlar: TensorFlow grafiğini, belirli bir model için özel olarak oluşturulan bir dizi hesaplama çekirdeğine derler. Bu çekirdekler modele özgü olduğundan optimizasyon için modele özgü bilgilerden yararlanabilirler. Örneğin, XLA'nın basit bir TensorFlow hesaplaması bağlamında yaptığı bir 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 üç çekirdek başlatır: biri çarpma, biri toplama ve biri çıkarma 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" 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". Birleştirme, XLA'nın en önemli optimizasyonudur. Bellek bant genişliği, donanım hızlandırıcılarda genellikle 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ştirme

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'sinde derlenmelidir semantiği vardır: İşlevin tamamı XLA ile derlenir veya errors.InvalidArgumentError istisnası atılır. XLA şu anda boyutların tahmin edilemediği işlevleri derleyemez. Yani tüm hesaplamanın tamamını çalıştırmadan tüm tensörlerin boyutlarını tahmin etmek mümkün değildir. Örneğin, aşağıdaki işlev derlenmez:

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

Ancak şekiller farklı koşular arasında değişiklik 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 Colab'ına, jit_compile=True kullanımıyla ilgili eğitim videosuna göz atı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ış stratejiyle kullanım

XLA:GPU, adım işlevine jit_compile=True 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üme oluşturma

TensorFlow modellerinde herhangi bir değişiklik yapmadan XLA'yı kullanmaya başlamanın basit bir yolu, TensorFlow işlevlerinde XLA kullanılarak derlenip çalıştırılabilen kümeler (bağlı alt grafikler) otomatik olarak bulan otomatik küme oluşturma özelliğini etkinleştirmektir. GPU'da otomatik küme oluşturma, 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üme oluşturma şu anda GPU iş yükleri için optimize edilmiştir ancak --tf_xla_cpu_global_jit işareti de 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üme oluşturma eğitim Colab'ına göz atın.

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

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

XLA, oluşturulan programları incelemenize olanak tanıyan iç gözlem olanakları sağlar. Oluşturulan programları dökmek 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

Yedekleme işlemi tamamlandıktan sonra /tmp/generated içinde aşağıdaki dosyaları bulabilirsiniz:

  • module_XXXX.*_optimizations.txt Her derlenmiş küme için bir tane olmak üzere oluşturulan XLA programları. XLA hata raporlarını gönderirken bu dosyaları eklemek çok faydalıdır.

  • module_XXXX.ir-*.ll NVPTX dahili işlevleriyle LLVM ara temsilinde oluşturulan dosyalar.

  • module_XXXX.ptx Oluşturulan PTX dosyaları.

Ayrıca, XLA kümelerinin TensorFlow grafiğine yerleştirilmesini görselleştiren grafiği şu yöntemlerle de dökebilirsiniz:

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

Yeniden oluşturulabilir hata raporları

Oluşturulan XLA programlarının ve kullanılan otomatik küme oluşturma yerleştirmesinin dökümlerini içeren bir hata raporunun yeniden oluşturulması çok daha kolaydır. Otomatik küme oluşturma ile çalışan bir TensorFlow programı için bunları oluşturmak üzere şunları 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"

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

Mümkünse run_hlo_module kullanarak ve oluşturulan programlarda iteratif olarak çalıştırarak hatayı tek bir XLA programına ayırmayı deneyin.

Daha fazla bilgi

XLA Ön Uçları

XLA programları TensorFlow'un yanı sıra aşağıdakiler tarafından da 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ığı

Sohbetler

jit_compile=True ile TF'den XLA kullanma

XLA'ya Genel Bakış