Kode error: 0100

Kategori: Runtime: Kegagalan Alokasi Buffer

Error ini menunjukkan bahwa pengalokasi memori runtime XLA:TPU gagal menemukan blok memori yang sesuai di HBM akselerator untuk alokasi yang diminta.

Contoh Pesan Error:

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

Backend XLA: TPU

Ringkasan

Error ini terjadi pada

  • kegagalan alokasi buffer yang dimulai pengguna melalui jax.device_put atau
  • kegagalan alokasi output program yang dijadwalkan pengguna.

Kegagalan ini biasanya disebabkan oleh beberapa alasan:

  • Kehabisan Memori (OOM): Pengguna mencoba mengalokasikan bagian memori yang lebih besar daripada jumlah total memori kosong yang tersedia di HBM TPU.
  • Fragmentasi Memori: Alokasi gagal karena tidak ada satu blok kosong yang berdekatan di ruang memori yang cukup besar untuk memenuhi ukuran yang diminta. Jumlah total memori kosong cukup untuk alokasi, tetapi tersebar di seluruh ruang memori dalam blok kecil yang tidak berdekatan.

Runtime TPU memiliki sejumlah mekanisme untuk mencoba ulang kegagalan alokasi, termasuk:

  • Jika ada pelepasan yang diantrekan, runtime akan mencoba kembali alokasi yang gagal,
  • Pada OOM yang disebabkan oleh fragmentasi, runtime dapat secara otomatis memicu defragmentasi dan percobaan ulang.
  • Runtime TPU memprioritaskan alokasi buffer daripada menjaga program tetap dimuat. Jika alokasi buffer gagal karena HBM tidak mencukupi, sistem akan mengeluarkan program TPU yang dimuat hingga memori yang cukup tersedia untuk buffer.

Jadi, error yang terjadi setelah mitigasi di atas biasanya memerlukan tindakan pengguna.

Proses debug

  • Mengurangi jejak memori model Anda:
    • Kurangi Ukuran Batch: Mengurangi ukuran batch secara langsung akan menurunkan penggunaan memori.
    • Penyusunan Parameter: Untuk model yang sangat besar, gunakan teknik seperti paralelisme model atau penyusunan untuk mendistribusikan parameter di seluruh HBM beberapa host atau core TPU.
    • Memperpendek Panjang Urutan/Konteks: Untuk model yang beroperasi pada urutan (seperti model bahasa), mengurangi panjang urutan input dapat mengurangi jejak memori secara signifikan.
    • Donasi Buffer: Memanfaatkan fitur framework (seperti: jax.jit(..., donate_argnums=...)) untuk memberi sinyal ke XLA bahwa buffer input tertentu dapat diganti dan digunakan kembali untuk output.
    • Mengoptimalkan Strategi Checkpoint: Daripada menyimpan seluruh status model sekaligus, sebaiknya simpan hanya bobot model atau gunakan strategi checkpointing yang di-shard.
  • Tata Letak dan Padding Memori Alamat:
    • Memori TPU dialokasikan dalam potongan, dan padding dapat meningkatkan ukuran tensor yang sebenarnya.
  • 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 membuat laporan mendetail tentang semua alokasi saat ini jika terjadi OOM, yang dapat sangat berharga untuk proses debug.