HLO Hesaplamalarını Dökme

HLO dökümü, hesaplamanın farklı aşamalarındaki HLO modüllerinin metin biçimindeki gösterimidir. Hata ayıklama için yararlıdır ve genellikle hata raporlarına eklemeniz gerekir. Bu, genellikle HLO talimatlarını ve özelliklerini listeleyen, kullanıcıların okuyabileceği bir metin dosyasıdır. Bazen HLO modülleri şu şekilde dökülür:

  • HloProto: Daha yapılandırılmış ve makine tarafından okunabilir bir biçim olan protokol arabellek dosyaları.
  • HloSnapshot: HLO modülü ve girişleri. HLO'ları yeniden oynatmak için bazen rastgele veriler yerine belirli bir hesaplamaya beslenen gerçek girişler gerekir.

Dökümleri belirtmek ve almak için XLA işaretlerini kullanabilirsiniz. Çoğu durumda, bunu bir ortam değişkeniyle ayarlayabilirsiniz. JAX, HLO dökümünü yazdırmak için programatik bir yöntem de sunar.

Yerel Yürütme

Ortam Değişkenlerini Kullanma

Dökümleri almak için XLA_FLAGS ortam değişkenini gerekli işaretlerle ayarlayabilirsiniz. Bu özellik JAX, TensorFlow ve PyTorch/XLA'da kullanılabilir.

HLO modüllerini ve diğer hata ayıklama bilgilerini belirli bir dizine boşaltmak için programınızı --xla_dump_to işaretiyle çalıştırın:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Örneğin, yollar olarak /tmp veya /tmp/xladump kullanabilirsiniz.

Bu, varsayılan olarak HLO modüllerini optimizasyon işlem hattının en başında ve sonunda metin olarak boşaltır.

Biçimi açıkça da belirtebilirsiniz:

  1. Metin dökümleri
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. HLO protokolleri
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. HLO anlık görüntüleri
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Graphviz sunucusuyla grafik oluşturma (yalnızca küçük grafiklerde iyi çalışır)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Grafiği HTML dosyasına dönüştürme (yalnızca küçük grafiklerde iyi çalışır)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Daha büyük grafiklerde, grafiğin bölümlerini görselleştirmek için interactive_graphviz simgesini kullanabilirsiniz.

Belirli ara geçişleri dökme

Önceden optimize edilmiş / nihai olarak optimize edilmiş standart HLO'lara ek olarak, belirli bir derleyici geçişinden sonra HLO'ların durumunu da dökebilirsiniz.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

HLO modülleri, adları normal ifadeyle (regex) eşleşen geçişler için dökülür. Örneğin, SPMD bölümleme ile ilgili geçişlerden kaynaklanan HLO'ları aşağıdaki gibi gözlemleyebilirsiniz:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

Her XLA geçişinden sonra sonucu boşaltmak için (bu işlem çok sayıda dosya oluşturur) şunları ayarlayabilirsiniz:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX'e özgü seçenekler

JAX'te programatik olarak

İşaret veya ortam değişkeni iletmek yerine, JAX'in lower ve compile API'lerini kullanarak HLO'yu programatik olarak da dökebilirsiniz.

Aşağıdaki komutla optimize edilmemiş orijinal HLO'yu yerel olarak getirin:

jax.jit(f).lower(*args).as_text('hlo')

HLO derleme geçişleri sırasında dosyalara döküm için şunu belirtin:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Dump jaxprs

jaxprs, program izleri için JAX'in ara gösterimidir. Bunu boşaltmak için ortam değişkenlerini ayarlayın:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Aşamalı olarak dışa aktarılan hesaplamaları dışa aktarma ve serileştirme: Hata ayıklama başlıklı JAX dokümanında daha fazla bilgi edinin.

Google Colab

Ortam değişkenleri

Not defterinizin ilk çalıştırılan hücresine (çünkü ortam değişkenleri ve komut satırı işaretleri genellikle yalnızca bir kez işlenir, örneğin modül içe aktarma sırasında veya XLA arka uç başlatma sırasında) yukarıda XLA_FLAGS ile ayrıntılı olarak açıklanan os.environ'yı ekleyin. Örneğin:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

Bu işlem, hesaplamayı DIRECTORY_PATH konumuna aktarır. Örneğin, /tmp. Bu dizini görüntülemek ve erişmek için Colab'de sol kenar çubuğundaki "Dosyalar" tarayıcısına gidin.

Yerel yürütme bölümünde belirtilen tüm işaretleri kullanabilirsiniz.

JAX'e özgü seçenekler

Yerel yürütmeye benzer şekilde, canlı ve etkileşimli iç gözlem için bir hesaplamanın önceden optimize edilmiş HLO'sunu doğrudan yazdırabilirsiniz:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

Ayrıca, bir hesaplamanın optimize edilmiş HLO'sunu doğrudan yazdırabilirsiniz:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Tüm/Küçük Hesaplamaları Dökme

Küçük derlemeler de dahil olmak üzere her şeyi bir dökümde görmek istiyorsanız JAX ortam değişkenini ayarlayın:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mozaik

Mosaic, Pallas TPU arka ucu ve deneysel Pallas GPU arka ucu için bir derleyicidir. Mozaik hesaplamayı boşaltmak için aşağıdaki işareti ayarlayın:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Alternatif olarak, TPU başlatma bağımsız değişkenlerini ortam değişkeni olarak ayarlayabilirsiniz:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Daha fazla bilgi edinmek için Pallas ve Mosaic ile ilgili JAX dokümanlarına göz atın.

HLO Dumps'ın yer aldığı daha fazla içerik

Doğru hesaplamayı bulma

Genellikle birçok hesaplama dökülür. Dökümü alınan dosyalar, günlüklerde belirtilen JAX, TensorFlow veya PyTorch/XLA "hesaplama adı" ile açıkça adlandırılır. Bu sayede ilgili HLO dosyalarını kolayca tanımlayabilirsiniz. Örneğin:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Aksi takdirde, belirli sembollerin veya hesaplamaların hangi modülde bulunduğunu hızlıca belirlemek için ripgrep simgesini kullanabilirsiniz.

İpucu: Hata raporlarınıza ilgilendiğiniz 3 döküm öncesi/sonrası/arabellek atama dosyasını ekleyin.

HLO Dönüşümü

HLOProto ve metin biçimleri arasında çeviri yapabilen hlo-opt adlı bir araç. Bir biçiminiz olduğu ancak hata ayıklama için diğer biçime ihtiyacınız olduğu durumlarda kullanışlıdır.

Kullanmayı öğrenin: XLA Tooling belgeleri: hlo-opt.

Tekrar Oynat

Dökümü alınmış hesaplamaları, sahte veriler veya giriş anlık görüntüleriyle belirli bir XLA arka ucunda çalıştırabilirsiniz (tekrar oynatabilirsiniz). Bu, XLA'daki sorunları yeniden oluşturmak, yinelemek ve hatalarını ayıklamak için kullanışlı bir yöntemdir.

Aşağıdaki komutlarda sahte veriler kullanılmaktadır. HLO anlık görüntülerini kaydettiyseniz bunları gönderebilirsiniz. Bu durumda, anlık görüntüdeki veriler kullanılır. Anlık görüntü çalıştırılırken sahte verileri kullanmaya devam etmek için --force_fake_data işaretini iletin.

CPU arka ucu:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

GPU arka ucu:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt