Kode error: E1001

Kategori: Waktu Kompilasi: Scoped Vmem OOM

Error ini menunjukkan bahwa program memerlukan lebih banyak Scoped Vector Memory (Vmem) daripada yang dialokasikan.

Contoh Pesan Error:

RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...

Backend XLA: TPU

Ringkasan

TPU memiliki Vector Memory (VMEM) yang merupakan memori scratchpad lokal yang digunakan secara eksklusif oleh TensorCore (TC). Compiler mengelola Vmem untuk berbagai jenis alokasi:

  • Alokasi cakupan petunjuk: Penyimpanan sementara di Vmem saat menjalankan satu petunjuk HLO. Hal ini mencakup buffer rentang operand (misalnya, untuk buffering ganda) dan tumpahan register.
  • Alokasi cakupan program: Alokasi yang berada di luar cakupan satu instruksi HLO. Ini biasanya merupakan hasil sementara dan perantara HLO yang merupakan input dan/atau output dari instruksi HLO.

Error Kehabisan Memori Vmem dengan Cakupan Waktu Kompilasi terjadi saat alokasi dengan cakupan petunjuk melampaui batas alokasi untuk petunjuk tersebut. Batas ini dikontrol

dan

Error ini biasanya disebabkan oleh bug compiler internal atau oleh kernel kustom yang melampaui batas alokasinya.

Proses debug

Analisis pesan error dengan cermat untuk mengidentifikasi apakah error berasal dari kernel kustom atau HLO standar. Error karena kernel kustom harus memiliki tanda tangan berikut:

Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
  • Vmem OOM dengan cakupan kernel kustom: Jika error mengarah ke kernel kustom → Lanjutkan ke Menyesuaikan Ulang Kernel.
  • Masalah Vmem non-kernel: Jika OOM Vmem terjadi karena operasi non-kernel kustom, kemungkinan besar ini adalah bug compiler internal. Buat laporan bug di XLA dengan dump HLO.

Menyesuaikan kembali kernel

Jika error berasal dari kernel kustom, gunakan teknik berikut untuk mengurangi persyaratan memori kernel:

  • Sesuaikan Ukuran Blok: Kurangi ukuran blok (ukuran petak) dalam konfigurasi kernel Anda, untuk menurunkan penggunaan Vmem yang Tercakup.
  • Menetapkan Batas Vmem yang Tercakup Per-Kernel: Minta secara eksplisit jumlah memori yang diperlukan untuk kernel tertentu menggunakan param vmem_limit_bytes
  • Mengubah Warna Memori: Mewarnai/membatasi input/output kernel secara eksplisit ke Vmem menggunakan pallas.tpu.with_memory_space_constraint. Berhati-hatilah agar tidak mewarnai terlalu banyak input output ke Vmem, karena dapat menyebabkan OOM Vmem secara keseluruhan.
  • Menyesuaikan batas Vmem: Jika penyesuaian ulang khusus kernel sulit dilakukan atau masalahnya memengaruhi banyak kernel, Anda dapat menyesuaikan batas Vmem global menggunakan tanda --xla_tpu_scoped_vmem_limit_kib.