Kode error: 0101

Kategori: Runtime: Kegagalan Alokasi Program

Error ini menunjukkan bahwa runtime XLA pada perangkat TPU gagal memuat file yang dapat dieksekusi program XLA yang dikompilasi ke dalam HBM TPU.

Contoh Pesan Error:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Backend XLA: TPU

Ringkasan

Error ini biasanya disebabkan oleh salah satu alasan berikut:

  • Ukuran Program Melebihi HBM yang Tersedia: Program XLA yang dikompilasi, termasuk petunjuk, data statis, dan konstanta tersemat, lebih besar daripada total jumlah HBM kosong yang saat ini tersedia di core TPU tertentu tempat program dimuat.
  • Fragmentasi HBM: Meskipun total HBM kosong di perangkat mungkin cukup secara keseluruhan, HBM tersebut tidak tersedia dalam satu blok berdekatan yang cukup besar untuk memuat seluruh program.

Anda harus memahami cara runtime TPU memprioritaskan memori. Alokasi Buffer mendapat hak istimewa atas program yang dimuat. Jika alokasi buffer gagal, runtime akan mengeluarkan program yang sudah dimuat dari HBM untuk mengosongkan ruang. Hal ini dapat menyebabkan program yang sebelumnya berhasil dimuat kini gagal dengan error OOM, karena HBM kini ditempati oleh lebih banyak buffer data.

Proses debug

  • Mengurangi Jejak Memori Buffer: Mengosongkan memori yang digunakan oleh buffer data akan menyisakan lebih banyak ruang untuk program itu sendiri:
    • Kurangi Ukuran Batch: Ini adalah salah satu cara paling efektif untuk mengurangi jumlah memori yang digunakan untuk aktivasi.
    • Parameter Sharding: Untuk model yang sangat besar, gunakan paralelisme model atau teknik sharding (seperti FSDP atau Megascale) untuk mendistribusikan parameter dan komputasi model di beberapa host atau core TPU.
    • Memperpendek Panjang Urutan/Konteks: Untuk model yang memproses data berurutan (misalnya, NLP), mengurangi panjang urutan dapat secara signifikan mengurangi penggunaan memori.
    • Donasi Buffer: Menggunakan fitur framework (misalnya, jax.jit(..., donate_argnums=...)) untuk memungkinkan XLA menggunakan kembali memori buffer input untuk menyimpan output, sehingga mengurangi penggunaan memori puncak.
  • Mengurangi persyaratan memori program untuk sementara:
    • Mengurangi penggunaan memori program untuk sementara dengan menggunakan flag tpu_shared_memory_percent. Perhatikan bahwa hal ini dapat memengaruhi performa secara negatif.
  • Mengoptimalkan Strategi Eksekusi/Mengurangi beban Penayangan:
    • Mengelola Pemuatan Program: Jika Anda mengompilasi JIT beberapa fungsi, perlu diketahui bahwa setiap fungsi dapat menyebabkan program dimuat. Coba susun workload Anda untuk meminimalkan jumlah program yang dimuat secara bersamaan.
  • Pastikan tidak ada kebocoran memori:
    • Pastikan referensi ke objek jax.Array tidak ditahan lebih lama dari yang dimaksudkan. Mempertahankan objek jax.Array dapat mencegah pelepasan alokasi otomatis meskipun kompilasi program telah selesai.

Alat

  • Aktifkan tanda tpu_log_allocations_on_oom yang alokatornya akan membuang laporan mendetail tentang semua alokasi saat ini jika terjadi OOM, yang dapat sangat berharga untuk proses debug.
  • Membuat Profil Program Anda: Gunakan profiler memori JAX atau profiler TensorFlow untuk mendapatkan tampilan mendetail tentang penggunaan memori program Anda dari waktu ke waktu. Hal ini dapat membantu mengidentifikasi lonjakan penggunaan memori yang tidak terduga.