Kode error: 1000

Kategori: Waktu Kompilasi: HBM OOM

Error ini menunjukkan bahwa program memerlukan High Bandwidth Memory (HBM) yang lebih besar daripada yang tersedia secara fisik di perangkat TPU.

Contoh Pesan Error:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

Backend XLA: TPU

Ringkasan

XLA melakukan pemeriksaan untuk memastikan ukuran gabungan semua alokasi statis yang diperlukan sesuai dengan HBM perangkat.

Compiler mengelola kapasitas HBM tetap TPU untuk beberapa jenis alokasi:

• Input dan Output Program: Batch pelatihan, status pengoptimal, dll.
• Temporer TensorCore + SparseCore: Memori dinamis yang diperlukan untuk kalkulasi perantara (misalnya, aktivasi, gradien, dll.).
• Biner yang Dikompilasi: Kode mesin untuk TensorCore (TC) dan SparseCore (SC).
• Overhead Sistem: Ruang yang dicadangkan untuk Runtime XLA (misalnya, buffer feed pada generasi TPU yang lebih lama).
• Konstanta: Nilai konstanta yang disematkan dalam IR HLO dialokasikan di HBM.
• Internal Compiler: Alokasi tingkat program dan per-HLO (misalnya, info perutean untuk node dalam mesh)

Error ini terjadi saat compiler XLA tidak dapat menyesuaikan semua alokasi di atas ke dalam HBM perangkat.

Proses debug

Analisis pesan error dan log dengan cermat untuk menentukan kategori HBM OOM di bawah yang paling sesuai dengan error Anda:

• Penggunaan HBM TensorCore (TC) + SparseCore (SC) Melebihi Batas: Jika error secara eksplisit menguraikan penggunaan, misalnya, "Penggunaan Hbm TC: X, Penggunaan Hbm SC: Y". → Langsung ke Bagian 1. Seimbangkan penggunaan HBM TC dan SC.
• Alokasi yang Sangat Besar dan Tidak Terduga: Jika errornya berbunyi "Ran out of memory in memory space HBM", periksa log untuk melihat enumerasi alokasi terbesar di HBM. Jika ada satu atau beberapa tensor besar yang tidak terduga (misalnya, > 50% dari batas HBM) → Lanjutkan ke Bagian 2. Alokasi yang Sangat Besar dan Tidak Terduga.
• Alokasi Gabungan Melebihi Batas HBM: Jika error berbunyi "Ran out of memory in memory space HBM", tetapi tidak ada tensor besar yang tidak terduga dalam log → Lanjutkan ke Bagian 3. Alokasi Gabungan Melebihi Batas HBM.

Bagian 1. Menyeimbangkan penggunaan HBM TC dan SC

Jika error secara eksplisit menguraikan penggunaan, misalnya, "Penggunaan Hbm TC: X, Penggunaan Hbm SC: Y" membandingkan kedua nilai untuk mengidentifikasi hambatan

• Penggunaan SparseCore Tinggi:
  • Mengoptimalkan Penggunaan Stack HBM: Penggunaan memori stack HBM diskalakan dengan feature_width, max_unique_nz_per_row, dan logical_replica_count. Anda dapat mengurangi penggunaan stack puncak dengan menyesuaikan flag --xla_sc_num_serialized_tables_to_optimize_hbm yang menserialkan pemrosesan tabel. Hal ini akan mengurangi paralelisme.
  • Periksa Overhead Padding: SparseCore menyelaraskan tabel penyematan ke 32B (8 float). Tabel dengan lebar fitur kecil (misalnya, < 8 float) menimbulkan overhead padding yang signifikan, sehingga membuang HBM.
  • Mengurangi Penggunaan Heap: Nilai tinggi untuk maximum_parallel_iterations meningkatkan jumlah data input yang telah diambil sebelumnya ke dalam heap HBM. Menurunkan nilai ini dapat membebaskan memori yang signifikan.
  • Verifikasi Sharding: Pastikan tabel penyematan di-shard dengan benar di semua chip. Lihat Cara batas diterjemahkan ke dalam tabel.
  • Lihat SC: Hambatan performa dan memori untuk mendapatkan ide lainnya.
• Penggunaan TensorCore Tinggi:
  • Lanjutkan ke Bagian 2.
• Seimbang
  • Jika masing-masing tidak berlebihan, tetapi jumlahnya terlalu tinggi, Anda berada pada kapasitas chip. Anda harus mencoba mengurangi penggunaan kedua komponen. Ikuti rekomendasi di ketiga bagian.

Bagian 2. Alokasi yang Tiba-Tiba Besar

Jika satu atau beberapa alokasi berukuran besar yang tidak terduga ada dalam log (> 50% batas HBM), hampir tidak pernah menjadi masalah kapasitas hardware. Biasanya ini adalah error konfigurasi. Periksa label XLA (jika ada) alokasi besar, untuk mendapatkan petunjuk tentang kode sumber JAX-nya.

• Menghapus Artefak Proses Debug:
  • Menggunakan jax.debug.print() dalam proses yang berskala besar dapat memaksa compiler untuk mewujudkan tensor penuh di HBM untuk mentransfernya ke CPU, sehingga menghentikan penggabungan dan meningkatkan penggunaan memori puncak. Hapus semua jax.debug.print() yang tersisa.
• Memperbaiki Bentuk atau Sharding Mesh yang Tidak Efisien:
  • Bentuk mesh yang salah atau anotasi sharding yang hilang dapat menyebabkan compiler secara default menggunakan Replication - sehingga compiler mencoba menyesuaikan tensor yang sangat besar pada satu chip
  • Periksa bentuk alokasi besar dan verifikasi bahwa sharding ditentukan dan dipropagasi dengan benar oleh XLA.

Bagian 3. Alokasi Gabungan Melebihi Batas HBM

Jika program kehabisan kapasitas karena jumlah total alokasi melebihi batas HBM, sering kali berguna untuk memvisualisasikan profil memori guna mengidentifikasi buffer tertentu yang berkontribusi pada penggunaan puncak. Lihat Men-debug error OOM dengan XProf untuk mendapatkan panduan langkah demi langkah tentang cara mengidentifikasi kontributor memori puncak.

Setelah Anda mengidentifikasi beberapa kontributor teratas, gunakan langkah-langkah berikut untuk mengoptimalkan jejak memori.

A. Memeriksa padding dan perataan tensor

Bentuk tensor yang tidak efisien adalah penyebab umum dan tersembunyi dari error kehabisan memori (OOM) di TPU. Untuk mendapatkan performa puncak di TPU, XLA mengisi dimensi tensor—biasanya ke kelipatan 128 untuk dimensi paling kecil dan 8 untuk dimensi kedua terkecil. Pengisihan ini memengaruhi array input dan tensor perantara (HLO sementara), yang berpotensi meningkatkan penggunaan memori secara signifikan, terutama dengan ukuran dimensi kecil. Lihat Tata Letak Array.

• Memeriksa bentuk buffer besar: (Di TPU v5 dengan tata letak default)
  • Mengarahkan kursor ke buffer di Xprof Memory Viewer akan memunculkan kartu detail buffer yang berisi detail buffer termasuk informasi padding.
  • Contoh: Bentuk (129, 1024) dapat diisi hingga (256, 1024), sehingga menyebabkan pemborosan memori hampir 50%.
  • Koreksi: Bentuk (128, 1024) tidak memerlukan padding dan menimbulkan pemborosan memori 0%.
• Menyelaraskan dimensi: Pastikan semua dimensi tensor besar (ukuran batch, dimensi penyematan, ukuran tersembunyi) adalah kelipatan 128.

B. Menyesuaikan konfigurasi

Anda sering kali dapat menyelesaikan masalah OOM dengan penyesuaian konfigurasi berikut:

• Kurangi Ukuran Batch: Memori yang diperlukan untuk aktivasi dan gradien perantara berbanding lurus dengan ukuran batch. Mengurangi ukuran batch sering kali dapat membantu mengurangi penggunaan memori.
• Menyumbangkan Buffer Input: Saat menggunakan jax.jit, tentukan donate_argnums untuk parameter model Anda. Hal ini memungkinkan XLA mengganti memori input dengan output.
• Aktifkan Presisi Campuran (bfloat16): Gunakan bfloat16 atau kuantisasi (int8, dll.) untuk tensor terbesar dalam program jika arsitektur model dan persyaratan kualitas memungkinkan.

C. Mengoptimalkan arsitektur dan sharding

Jika perubahan konfigurasi tidak memadai, topologi model mungkin terlalu besar untuk penyiapan hardware saat ini.

• Gunakan Generasi TPU yang Lebih Baru: TPU yang lebih baru umumnya menawarkan lebih banyak HBM per chip; beralihlah ke generasi TPU yang lebih baru jika tersedia.
• Jalankan pada topologi chip yang lebih besar: Jika bobot model terlalu besar untuk topologi yang ada, Anda dapat mencoba membaginya di lebih banyak chip.
• Menerapkan teknik sharding lanjutan:
  • Pelajari pendekatan paralelisme data, tensor, atau pipeline yang lebih canggih.
  • Tentukan petunjuk sharding untuk nilai dan output antara.
• Gunakan Pelepasan Host JAX: Lepaskan tensor besar ke memori CPU host. Misalnya, pelepasan aktivasi dan pelepasan status pengoptimal.

D. Menyesuaikan tanda XLA yang memengaruhi memori utama:

Flag memori utama dapat disetel untuk menyeimbangkan performa dengan penggunaan memori yang lebih rendah. Namun, tindakan ini harus digunakan sebagai upaya terakhir karena dapat memengaruhi performa secara negatif.

E. Menyesuaikan Pass Rematerialisasi XLA / Checkpointing Manual

Jika model hampir sesuai dengan memori, Anda dapat memaksa penerusan XLA::Rematerialization untuk memprioritaskan penghematan memori, yang berpotensi mengorbankan kompilasi yang lebih lambat:

Atau, gunakan dekorator jax.checkpoint dengan jax.grad untuk mengontrol secara manual perantara mana yang disimpan pada pass terus versus dihitung ulang pada pass mundur, dengan menukar siklus komputasi untuk HBM.

F. Menggunakan alat pembuatan profil lanjutan

Men-debug error OOM dengan XProf memberikan tutorial tentang penggunaan XProf Memory Viewer untuk memvisualisasikan penggunaan HBM dari sudut pandang compiler.

Alat ini memungkinkan Anda melihat alokasi memori puncak dan masa aktif buffer, yang sangat penting untuk memahami secara persis apa yang menggunakan HBM pada titik penggunaan puncak. Untuk penyiapan pembuatan profil umum, lihat Mulai menggunakan Xprof dan Pembuatan Profil TensorBoard.