Kode error: E1000

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: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

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.

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.
• TPU temporaries: Memori dinamis yang diperlukan untuk perhitungan menengah (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 infeed 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 TC: X, Penggunaan Hbm SC Y": Jika error secara eksplisit menguraikan penggunaan HBM, penggunaan gabungan TensorCore (TC) + SparseCore (SC) melampaui batas HBM. → Langsung ke Skenario 1. Menyeimbangkan penggunaan HBM TC dan SC.
• "Ran out of memory in memory space HBM": Periksa log untuk melihat enumerasi alokasi terbesar di HBM.
  • Jika ada satu atau beberapa tensor yang berukuran sangat besar secara tidak terduga (misalnya, > 50% dari batas HBM) → Lanjutkan ke Skenario 2. Kehabisan memori karena alokasi yang sangat besar secara tidak terduga.
  • Jika tidak ada tensor berukuran besar yang tidak terduga dalam log → Lanjutkan ke Skenario 3. Kehabisan memori karena alokasi gabungan.

Skenario 1. Menyeimbangkan penggunaan HBM TC dan SC

Jika error secara eksplisit menguraikan penggunaan, misalnya, "Penggunaan Hbm TC: X, Penggunaan Hbm SC: Y", berarti penggunaan gabungan TensorCore (TC) + SparseCore (SC) melampaui batas HBM. Bandingkan kedua nilai untuk mengidentifikasi hambatan:

• Penggunaan SparseCore yang tinggi
  • Mengoptimalkan penggunaan stack HBM: Konsumsi memori stack HBM diskalakan dengan feature_width, max_unique_nz_per_row, dan logical_replica_count. Anda dapat mengurangi penggunaan stack puncak dengan menyesuaikan tanda --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) akan menimbulkan overhead padding yang signifikan, sehingga membuang-buang HBM.
  • Mengurangi penggunaan heap: Nilai tinggi untuk maximum_parallel_iterations meningkatkan jumlah data input yang telah di-prefetch 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 yang tinggi
  • Lanjutkan ke Skenario 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.

Skenario 2. Kehabisan memori karena alokasi yang sangat besar secara tidak terduga

Jika Anda melihat pesan error "Ran out of memory in memory space HBM" dan satu atau beberapa alokasi yang sangat besar secara tidak terduga ada di log (> 50% batas HBM), hampir tidak pernah ada masalah kapasitas hardware. Error ini biasanya merupakan 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 pembagian atau bentuk mesh yang tidak efisien
  • Bentuk mesh yang salah atau anotasi sharding yang tidak ada dapat menyebabkan compiler secara default menggunakan replikasi - sehingga compiler mencoba menempatkan tensor yang sangat besar pada satu chip.
  • Periksa bentuk alokasi besar dan verifikasi bahwa penyiapan shard ditentukan dan dipropagasi dengan benar oleh XLA.

Skenario 3. Kehabisan memori karena alokasi gabungan

Jika Anda melihat pesan error "Ran out of memory in memory space HBM" dan tidak ada tensor besar yang tidak terduga dalam log, program kehabisan kapasitas karena jumlah total alokasi melebihi batas HBM. Dalam kasus ini, 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.

Skenario 3.A 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. Perhatikan bahwa perubahan ini dapat memengaruhi perilaku model dan harus dipertimbangkan dengan cermat.

Skenario 3.B 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 membagi bobot tersebut 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: Teknik pelepasan host memungkinkan pengguna melepas tensor besar ke memori CPU host (misalnya, pelepasan aktivasi dan pelepasan status pengoptimal).

Skenario 3.C Memeriksa padding dan perataan tensor

Bentuk tensor yang tidak efisien adalah penyebab umum OOM yang tidak terlihat 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 menampilkan kartu detail buffer yang berisi detail buffer termasuk informasi padding.
  • Contoh: Bentuk (129, 1024) dapat diisi dengan (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. Perhatikan bahwa perubahan ini dapat memengaruhi perilaku model dan harus dipertimbangkan dengan cermat.

Skenario 3.D Menyesuaikan flag XLA yang memengaruhi memori utama

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

Skenario 3.E Penyesuaian XLA rematerialization pass/checkpointing manual

Jika model hampir sesuai dengan memori, Anda dapat menggunakan decorator 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.

Atau, Anda dapat memaksa XLA::Rematerialization untuk memprioritaskan penghematan memori, yang berpotensi mengorbankan kompilasi yang lebih lambat:

Perhatikan bahwa perubahan pada tanda XLA harus digunakan sebagai langkah terakhir, karena dapat memengaruhi performa secara negatif.

Skenario 3.F Menggunakan alat pembuatan profil lanjutan

Men-debug error OOM dengan XProf memberikan tutorial tentang penggunaan XProf Memory Viewer untuk memvisualisasikan penggunaan HBM dari tampilan 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.