Men-debug error OOM dengan XProf

Error Kehabisan Memori (OOM) terjadi saat kapasitas Memori Bandwidth Tinggi (HBM) akselerator (GPU atau TPU) habis. Beberapa penyebab umum masalah OOM dan teknik pen-debug-an dijelaskan secara mendetail dalam dokumentasi E1000 - Compile Time HBM OOM dan dokumentasi JAX tentang alokasi memori GPU.

Halaman ini menjelaskan cara menggunakan alat Memory Viewer XProf untuk memvisualisasikan penggunaan memori program JAX Anda, mengidentifikasi instance penggunaan puncak, dan men-debug error OOM. Hal ini melibatkan langkah-langkah berikut:

  1. Jalankan program Anda dengan jax.profiler.trace untuk merekam profil.
  2. Mulai XProf di latar belakang, dan gunakan alat Memory Viewer untuk melihat detail penggunaan memori.

Contoh program

Program JAX berikut menyebabkan error OOM:

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

Pada mesin TPU, program ini gagal dengan:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(Di mesin GPU, error akan terlihat seperti: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

Menjalankan XProf

Instal xprof (pip install xprof), dan mulai instance XProf dengan menentukan direktori tempat profil disimpan:

xprof --logdir=/tmp/xprof/ --port=6006

Buka instance (di mesin lokal, di http://localhost:6006). Di dropdown Tools, pilih Memory Viewer, lalu di jendela alat Memory Viewer, pilih HBM di dropdown Memory Types (biasanya dipilih secara default).

Halaman XProf Memory Viewer untuk contoh program di atas

Dokumentasi alat XProf: Memory Viewer menjelaskan komponen alat dan informasi yang ditampilkan.

Fokus pada bagian Operasi HLO pada Alokasi Memori Puncak yang menampilkan tiga diagram buffer pada titik penggunaan memori puncak. Buffer mencakup: * Input dan Output Program: Batch pelatihan, status pengoptimal, dll. * Sementara TensorCore dan SparseCore: Memori dinamis yang diperlukan untuk perhitungan perantara (seperti aktivasi, gradien, dll.)

Anda dapat mengarahkan kursor ke diagram buffer untuk mendapatkan detail selengkapnya tentang Op seperti ukuran, bentuk, jenis alokasi, dan lainnya. Hal ini dapat membantu Anda mengidentifikasi dan mengevaluasi operasi yang mungkin memiliki sementara yang tinggi atau berdurasi lama, tensor input/perantara/output besar yang memiliki padding tidak efisien, dll., yang berkontribusi pada memori puncak dan perlu disesuaikan atau dioptimalkan.

Pelajari teknik proses debug tertentu di E1000: Proses Debug.