Membuat profil komputasi JAX dengan XProf

XProf adalah cara yang bagus untuk mendapatkan dan memvisualisasikan rekaman aktivitas dan profil performa program Anda, termasuk aktivitas di GPU dan TPU. Hasil akhirnya akan terlihat seperti ini:

Contoh XProf

Pengambilan gambar terprogram

Anda dapat menginstrumentasikan kode untuk merekam rekaman aktivitas profiler untuk kode JAX melalui metode jax.profiler.start_trace dan jax.profiler.stop_trace. Panggil jax.profiler.start_trace dengan direktori untuk menulis file rekaman aktivitas. Direktori ini harus sama dengan direktori --logdir yang digunakan untuk memulai XProf. Kemudian, Anda dapat menggunakan XProf untuk melihat rekaman aktivitas.

Misalnya, untuk merekam aktivitas profiler:

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()

Perhatikan panggilan jax.block_until_ready. Kami menggunakan ini untuk memastikan eksekusi di perangkat direkam oleh rekaman aktivitas. Lihat Pengiriman asinkron untuk mengetahui detail mengapa hal ini diperlukan.

Anda juga dapat menggunakan pengelola konteks jax.profiler.trace sebagai alternatif untuk start_trace dan stop_trace:

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()

Melihat rekaman aktivitas

Setelah merekam aktivitas, Anda dapat melihatnya menggunakan UI XProf.

Anda dapat meluncurkan UI profiler secara langsung menggunakan perintah XProf mandiri dengan mengarahkan perintah tersebut ke direktori log Anda:

$ 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)

Buka URL yang disediakan (misalnya, http://localhost:8791/) di browser Anda untuk melihat profil.

Rekaman aktivitas yang tersedia akan muncul di menu dropdown "Sesi" di sebelah kiri. Pilih sesi yang Anda minati, lalu di bagian dropdown "Alat", pilih "Trace Viewer". Sekarang Anda akan melihat linimasa eksekusi. Anda dapat menggunakan tombol WASD untuk menavigasi rekaman aktivitas, dan mengklik atau menarik untuk memilih peristiwa guna melihat detail selengkapnya. Lihat dokumentasi Alat Trace Viewer untuk mengetahui detail selengkapnya tentang cara menggunakan Trace Viewer.

Pengambilan manual melalui XProf

Berikut adalah petunjuk untuk merekam rekaman aktivitas N detik yang dipicu secara manual dari program yang sedang berjalan.

  1. Mulai server XProf:

    xprof --logdir /tmp/profile-data/

    Anda akan dapat memuat XProf di <http://localhost:8791/>. Anda dapat menentukan port yang berbeda dengan tanda --port.

  2. Di program atau proses Python yang ingin Anda buat profilnya, tambahkan kode berikut di dekat awal:

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

    Tindakan ini akan memulai server profiler yang terhubung ke XProf. Server profiler harus berjalan sebelum Anda melanjutkan ke langkah berikutnya. Setelah selesai menggunakan server, Anda dapat memanggil jax.profiler.stop_server() untuk mematikannya.

    Jika Anda ingin memprofilkan cuplikan program yang berjalan lama (misalnya, loop pelatihan yang panjang), Anda dapat menempatkan ini di awal program dan memulai program seperti biasa. Jika Anda ingin membuat profil program singkat (misalnya, microbenchmark), salah satu opsinya adalah memulai server profiler di shell IPython, dan menjalankan program singkat dengan %run setelah memulai pengambilan di langkah berikutnya. Opsi lainnya adalah memulai server profiler di awal program dan menggunakan time.sleep() untuk memberi Anda waktu yang cukup untuk memulai pengambilan.

  3. Buka <http://localhost:8791/>, lalu klik tombol "REKAM PROFIL" di kiri atas. Masukkan "localhost:9999" sebagai URL layanan profil (ini adalah alamat server profiler yang Anda mulai pada langkah sebelumnya). Masukkan jumlah milidetik yang ingin Anda profilkan, lalu klik "CAPTURE".

  4. Jika kode yang ingin Anda profil belum berjalan (misalnya, jika Anda memulai server profiler di shell Python), jalankan kode tersebut saat pengambilan sedang berjalan.

  5. Setelah pengambilan selesai, XProf akan otomatis dimuat ulang. (Tidak semua fitur pembuatan profil XProf terhubung dengan JAX, sehingga pada awalnya mungkin tampak seperti tidak ada yang direkam.) Di sebelah kiri pada bagian "Alat", pilih "Trace Viewer".

Sekarang Anda akan melihat linimasa eksekusi. Anda dapat menggunakan tombol WASD untuk menavigasi rekaman aktivitas, dan mengklik atau menarik untuk memilih peristiwa guna melihat detail selengkapnya di bagian bawah. Lihat dokumentasi Alat Trace Viewer untuk mengetahui detail selengkapnya tentang penggunaan Trace Viewer.

XProf dan Tensorboard

XProf adalah alat dasar yang mendukung fungsi pembuatan profil dan pengambilan rekaman aktivitas di Tensorboard. Selama xprof diinstal, tab "Profil" akan ada di dalam Tensorboard. Penggunaan ini sama dengan meluncurkan XProf secara independen, selama diluncurkan dengan mengarah ke direktori log yang sama. Hal ini mencakup fungsi pengambilan, analisis, dan penayangan profil. XProf menggantikan fungsi tensorboard_plugin_profile yang sebelumnya direkomendasikan.

$ 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)

Menambahkan peristiwa rekaman aktivitas kustom

Secara default, peristiwa di pemeriksa rekaman aktivitas sebagian besar adalah fungsi JAX internal tingkat rendah. Anda dapat menambahkan peristiwa dan fungsi Anda sendiri menggunakan jax.profiler.TraceAnnotation dan jax.profiler.annotate_function dalam kode Anda.

Mengonfigurasi opsi profiler

Metode start_trace menerima parameter profiler_options opsional, yang memungkinkan kontrol terperinci atas perilaku profiler. Parameter ini harus berupa instance jax.profiler.ProfileOptions.

Misalnya, untuk menonaktifkan semua rekaman aktivitas python dan host:

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()

Opsi umum

  1. host_tracer_level: Menetapkan tingkat rekaman aktivitas untuk aktivitas sisi host.

    Nilai yang Didukung:

    • 0: Menonaktifkan sepenuhnya pelacakan host (CPU).
    • 1: Mengaktifkan pelacakan hanya untuk peristiwa TraceMe yang diinstrumentasikan pengguna.
    • 2: Mencakup rekaman aktivitas level 1 serta detail eksekusi program tingkat tinggi seperti operasi XLA yang mahal (default).
    • 3: Mencakup rekaman aktivitas level 2 serta detail eksekusi program tingkat rendah yang lebih verbose, seperti operasi XLA yang murah.

  2. device_tracer_level: Mengontrol apakah pelacakan perangkat diaktifkan.

    Nilai yang Didukung:

    • 0: Menonaktifkan pelacakan perangkat.
    • 1: Mengaktifkan pelacakan perangkat (default).

  3. python_tracer_level: Mengontrol apakah pelacakan Python diaktifkan.

    Nilai yang Didukung:

    • 0: Menonaktifkan pelacakan panggilan fungsi Python (default).
    • 1: Mengaktifkan pelacakan Python.

Opsi konfigurasi lanjutan

Opsi TPU

  1. tpu_trace_mode: Menentukan mode untuk pelacakan TPU.

    Nilai yang Didukung:

    • TRACE_ONLY_HOST: Artinya, hanya aktivitas sisi host (CPU) yang direkam, dan tidak ada rekaman aktivitas perangkat (TPU/GPU) yang dikumpulkan.
    • TRACE_ONLY_XLA: Artinya, hanya operasi tingkat XLA di perangkat yang dilacak.
    • TRACE_COMPUTE: Opsi ini melacak operasi komputasi di perangkat.
    • TRACE_COMPUTE_AND_SYNC: Opsi ini melacak operasi komputasi dan peristiwa sinkronisasi di perangkat.

    Jika "tpu_trace_mode" tidak disediakan, trace_mode akan menggunakan default TRACE_ONLY_XLA.

  2. tpu_num_sparse_cores_to_trace: Menentukan jumlah core jarang yang akan dilacak di TPU.

  3. tpu_num_sparse_core_tiles_to_trace: Menentukan jumlah petak dalam setiap inti jarang yang akan dilacak di TPU.

  4. tpu_num_chips_to_profile_per_task: Menentukan jumlah TPU chip yang akan diprofilkan per tugas.

Opsi GPU

Opsi berikut tersedia untuk pembuatan profil GPU:

  • gpu_max_callback_api_events: Menetapkan jumlah maksimum peristiwa yang dikumpulkan oleh CUPTI callback API. Nilai defaultnya adalah 2*1024*1024.
  • gpu_max_activity_api_events: Menetapkan jumlah maksimum peristiwa yang dikumpulkan oleh CUPTI Activity API. Nilai defaultnya adalah 2*1024*1024.
  • gpu_max_annotation_strings: Menetapkan jumlah maksimum string anotasi yang dapat dikumpulkan. Nilai defaultnya adalah 1024*1024.
  • gpu_enable_nvtx_tracking: Mengaktifkan pelacakan NVTX di CUPTI. Nilai defaultnya adalah False.
  • gpu_enable_cupti_activity_graph_trace: Mengaktifkan pelacakan grafik aktivitas CUPTI untuk grafik CUDA. Nilai defaultnya adalah False.
  • gpu_pm_sample_counters: String yang dipisahkan koma dari metrik Pemantauan Performa GPU yang akan dikumpulkan menggunakan fitur pengambilan sampel PM CUPTI (misalnya, "sm__cycles_active.avg.pct_of_peak_sustained_elapsed"). Pengambilan sampel PM dinonaktifkan secara default. Untuk metrik yang tersedia, lihat dokumentasi CUPTI NVIDIA.
  • gpu_pm_sample_interval_us: Menetapkan interval pengambilan sampel dalam mikrodetik untuk pengambilan sampel PM CUPTI. Nilai defaultnya adalah 500.
  • gpu_pm_sample_buffer_size_per_gpu_mb: Menetapkan ukuran buffer memori sistem per perangkat dalam MB untuk pengambilan sampel PM CUPTI. Defaultnya adalah 64 MB. Nilai maksimum yang didukung adalah 4 GB.
  • gpu_num_chips_to_profile_per_task: Menentukan jumlah perangkat GPU yang akan diprofilkan per tugas. Jika tidak ditentukan, ditetapkan ke 0, atau ditetapkan ke nilai yang tidak valid, semua GPU yang tersedia akan di-profil. Opsi ini dapat digunakan untuk mengurangi ukuran pengumpulan rekaman aktivitas.
  • gpu_dump_graph_node_mapping: Jika diaktifkan, akan mencatat informasi pemetaan node grafik CUDA ke dalam rekaman aktivitas. Nilai defaultnya adalah False.

Contoh:

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

Menampilkan InvalidArgumentError jika ada kunci atau nilai opsi yang tidak dikenal.