XProf ile JAX hesaplamalarında profil oluşturma

XProf, GPU ve TPU'daki etkinlik de dahil olmak üzere programınızın performans izlerini ve profillerini edinip görselleştirmenin harika bir yoludur. Sonuç aşağıdaki gibi görünür:

XProf örneği

Programatik yakalama

jax.profiler.start_trace ve jax.profiler.stop_trace yöntemlerini kullanarak JAX kodu için bir profil oluşturucu izi yakalamak üzere kodunuzu izleyebilirsiniz. İzleme dosyalarının yazılacağı dizinle birlikte Call jax.profiler.start_trace işlevini kullanın. Bu, XProf'u başlatmak için kullanılan --logdir diziniyle aynı olmalıdır. Ardından, izlemeleri görüntülemek için XProf'u kullanabilirsiniz.

Örneğin, bir profil oluşturucu izi almak için:

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

jax.block_until_ready numaralı aramayı not edin. Cihaz üzerinde yürütmenin izlemeye dahil edildiğinden emin olmak için bu bilgiyi kullanırız. Bunun neden gerekli olduğuyla ilgili ayrıntılar için Asenkron gönderme başlıklı makaleyi inceleyin.

start_trace ve stop_trace yerine alternatif olarak jax.profiler.trace bağlam yöneticisini de kullanabilirsiniz:

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

İzi görüntüleme

İz yakaladıktan sonra XProf kullanıcı arayüzünü kullanarak görüntüleyebilirsiniz.

Günlük dizininize yönlendirerek bağımsız XProf komutunu kullanarak profiler kullanıcı arayüzünü doğrudan başlatabilirsiniz:

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

Belirtilen URL'ye gidin (ör. http://localhost:8791/) tarayıcınızda görüntüleyin.

Kullanılabilir izler soldaki "Oturumlar" açılır menüsünde görünür. İlgilendiğiniz oturumu seçin ve "Araçlar" açılır listesinde "İzleyiciyi İzle"yi seçin. Artık yürütme zaman çizelgesini görebilirsiniz. İzde gezinmek için WASD tuşlarını kullanabilir, daha fazla ayrıntı için etkinlikleri seçmek üzere tıklayabilir veya sürükleyebilirsiniz. İzleyiciyi kullanma hakkında daha fazla bilgi için İzleyici Aracı dokümanlarına bakın.

XProf ile manuel yakalama

Aşağıda, çalışan bir programdan manuel olarak tetiklenen N saniyelik bir izlemenin nasıl yakalanacağıyla ilgili talimatlar verilmiştir.

  1. XProf sunucusunu başlatma:

    xprof --logdir /tmp/profile-data/

    XProf'u <http://localhost:8791/> adresinden yükleyebilirsiniz. --port işaretini kullanarak farklı bir bağlantı noktası belirtebilirsiniz.

  2. Profilini oluşturmak istediğiniz Python programına veya sürecine, başlangıca yakın bir yere aşağıdaki kodu ekleyin:

    import jax.profiler
jax.profiler.start_server(9999)

    Bu komut, XProf'un bağlandığı profil oluşturucu sunucusunu başlatır. Bir sonraki adıma geçmeden önce profiler sunucusunun çalışıyor olması gerekir. Sunucuyu kullanmayı bitirdiğinizde jax.profiler.stop_server() numaralı telefonu arayarak kapatabilirsiniz.

    Uzun süren bir programın (ör. uzun bir eğitim döngüsü) bir bölümünün profilini oluşturmak istiyorsanız bu kodu programın başına yerleştirebilir ve programınızı her zamanki gibi başlatabilirsiniz. Kısa bir programın (ör. mikro karşılaştırma) profilini oluşturmak istiyorsanız bir seçenek de profil oluşturucu sunucusunu bir IPython kabuğunda başlatmak ve sonraki adımda yakalama başlatıldıktan sonra kısa programı %run ile çalıştırmaktır. Diğer bir seçenek de programın başında profil oluşturucu sunucusunu başlatmak ve yakalamayı başlatmak için yeterli zamanı tanımak üzere time.sleep() kullanmaktır.

  3. <http://localhost:8791/> uygulamasını açın ve sol üstteki "PROFİLİ YAKALA" düğmesini tıklayın. Profil hizmeti URL'si olarak "localhost:9999" değerini girin (bu, önceki adımda başlattığınız profiler sunucusunun adresidir). Profil oluşturmak istediğiniz milisaniye sayısını girin ve "CAPTURE" (YAKALA) seçeneğini tıklayın.

  4. Profilini oluşturmak istediğiniz kod henüz çalışmıyorsa (ör. profil oluşturucu sunucusunu bir Python kabuğunda başlattıysanız) yakalama işlemi devam ederken kodu çalıştırın.

  5. Yakalamanın tamamlanmasının ardından XProf otomatik olarak yenilenmelidir. (XProf profil oluşturma özelliklerinin tümü JAX'e bağlı değildir. Bu nedenle, başlangıçta hiçbir şey yakalanmamış gibi görünebilir.) Sol taraftaki "Araçlar" bölümünde "İz Görüntüleyici"yi seçin.

Artık yürütme zaman çizelgesini görebilirsiniz. İzlemeye gitmek için WASD tuşlarını kullanabilir, daha fazla ayrıntı görmek için etkinlikleri seçmek üzere tıklayabilir veya sürükleyebilirsiniz. İzleyiciyi kullanma hakkında daha fazla bilgi için İzleyici Aracı dokümanlarına bakın.

XProf ve TensorBoard

XProf, TensorBoard'daki profilleme ve iz yakalama işlevini destekleyen temel araçtır. xprof yüklü olduğu sürece Tensorboard'da bir "Profil" sekmesi bulunur. Aynı günlük dizinine yönlendirilerek başlatıldığı sürece bu aracı kullanmak, XProf'u bağımsız olarak başlatmakla aynıdır. Buna profil yakalama, analiz ve görüntüleme işlevleri dahildir. XProf, daha önce önerilen tensorboard_plugin_profile işlevinin yerini alıyor.

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

Özel izleme etkinlikleri ekleme

İzleme görüntüleyicideki etkinlikler varsayılan olarak çoğunlukla düşük düzeyli dahili JAX işlevleridir. Kodunuzda jax.profiler.TraceAnnotation ve jax.profiler.annotate_function kullanarak kendi etkinliklerinizi ve işlevlerinizi ekleyebilirsiniz.

Profiler seçeneklerini yapılandırma

start_trace yöntemi, profil oluşturucunun davranışı üzerinde ayrıntılı kontrol sağlayan isteğe bağlı bir profiler_options parametresini kabul eder. Bu parametre, jax.profiler.ProfileOptions örneği olmalıdır.

Örneğin, tüm Python ve ana makine izlemelerini devre dışı bırakmak için:

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Genel seçenekler

  1. host_tracer_level: Ana makine tarafındaki etkinlikler için izleme düzeyini ayarlar.

    Desteklenen Değerler:

    • 0: Ana makine (CPU) izlemeyi tamamen devre dışı bırakır.
    • 1: Yalnızca kullanıcı tarafından oluşturulan TraceMe etkinliklerinin izlenmesini sağlar.
    • 2: 1. düzey izlemelerin yanı sıra pahalı XLA işlemleri gibi üst düzey program yürütme ayrıntılarını içerir (varsayılan).
    • 3: 2. düzey izlerin yanı sıra ucuz XLA işlemleri gibi daha ayrıntılı ve düşük düzeyli program yürütme ayrıntılarını içerir.

  2. device_tracer_level: Cihaz izlemenin etkinleştirilip etkinleştirilmeyeceğini kontrol eder.

    Desteklenen Değerler:

    • 0: Cihaz izlemeyi devre dışı bırakır.
    • 1: Cihaz izlemeyi etkinleştirir (varsayılan).

  3. python_tracer_level: Python izlemenin etkinleştirilip etkinleştirilmediğini kontrol eder.

    Desteklenen Değerler:

    • 0: Python işlev çağrısı izlemeyi devre dışı bırakır (varsayılan).
    • 1: Python izlemeyi etkinleştirir.

Gelişmiş yapılandırma seçenekleri

TPU seçenekleri

  1. tpu_trace_mode: TPU izleme modunu belirtir.

    Desteklenen Değerler:

    • TRACE_ONLY_HOST: Bu, yalnızca ana makine tarafındaki (CPU) etkinliklerin izlendiği ve cihaz (TPU/GPU) izlerinin toplanmadığı anlamına gelir.
    • TRACE_ONLY_XLA: Bu, cihazda yalnızca XLA düzeyindeki işlemlerin izlendiği anlamına gelir.
    • TRACE_COMPUTE: Bu, cihazdaki işlem işlemlerini izler.
    • TRACE_COMPUTE_AND_SYNC: Bu seçenek, cihazdaki hem bilgi işlem işlemlerini hem de senkronizasyon etkinliklerini izler.

    "tpu_trace_mode" sağlanmazsa trace_mode varsayılan olarak TRACE_ONLY_XLA olur.

  2. tpu_num_sparse_cores_to_trace: TPU'da izlenecek seyrek çekirdek sayısını belirtir.

  3. tpu_num_sparse_core_tiles_to_trace: TPU'da izlenecek her seyrek çekirdekteki döşeme sayısını belirtir.

  4. tpu_num_chips_to_profile_per_task: Görev başına profili oluşturulacak TPU çip sayısını belirtir.

GPU seçenekleri

GPU profil oluşturma için aşağıdaki seçenekler kullanılabilir:

  • gpu_max_callback_api_events: CUPTI geri çağırma API'si tarafından toplanan maksimum etkinlik sayısını ayarlar. Varsayılan olarak 2*1024*1024 değerine ayarlanır.
  • gpu_max_activity_api_events: CUPTI etkinlik API'si tarafından toplanan maksimum etkinlik sayısını ayarlar. Varsayılan olarak 2*1024*1024 değerine ayarlanır.
  • gpu_max_annotation_strings: Toplanabilecek maksimum açıklama dizesi sayısını ayarlar. Varsayılan olarak 1024*1024 değerine ayarlanır.
  • gpu_enable_nvtx_tracking: CUPTI'de NVTX izlemeyi etkinleştirir. Varsayılan olarak False değerine ayarlanır.
  • gpu_enable_cupti_activity_graph_trace: CUDA grafikleri için CUPTI etkinlik grafiği izlemeyi etkinleştirir. Varsayılan olarak False değerine ayarlanır.
  • gpu_pm_sample_counters: CUPTI'nin PM örnekleme özelliği kullanılarak toplanacak, virgülle ayrılmış GPU Performans İzleme metrikleri dizesi (ör. "sm__cycles_active.avg.pct_of_peak_sustained_elapsed"). PM örnekleme varsayılan olarak devre dışıdır. Kullanılabilen metrikler için NVIDIA'nın CUPTI belgelerine bakın.
  • gpu_pm_sample_interval_us: CUPTI PM örneklemesi için örnekleme aralığını mikrosaniye cinsinden ayarlar. Varsayılan olarak 500 değerine ayarlanır.
  • gpu_pm_sample_buffer_size_per_gpu_mb: CUPTI PM örneklemesi için cihaz başına sistem belleği arabellek boyutunu MB cinsinden ayarlar. Varsayılan olarak 64 MB değerine ayarlanır. Desteklenen maksimum değer 4 GB'tır.
  • gpu_num_chips_to_profile_per_task: Görev başına profili oluşturulacak GPU cihazlarının sayısını belirtir. Belirtilmezse, 0 olarak ayarlanırsa veya geçersiz bir değere ayarlanırsa kullanılabilir tüm GPU'lar profillendirilir. Bu, iz toplama boyutunu küçültmek için kullanılabilir.
  • gpu_dump_graph_node_mapping: Etkinleştirilirse CUDA grafiği düğümü eşleme bilgilerini izlemeye aktarır. Varsayılan olarak False değerine ayarlanır.

Örneğin:

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

Tanınmayan anahtarlar veya seçenek değerleri bulunursa InvalidArgumentError değerini döndürür.