Kategori: Waktu Kompilasi: Kegagalan Alokasi SparseCore
Error ini terjadi saat compiler XLA:SparseCore tidak dapat mengalokasikan blok memori yang berdekatan di ruang memori yang ditentukan yang diperlukan oleh program SparseCore saat ini.
Contoh pesan error:
INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>
Backend XLA: TPU
Ringkasan
SparseCore (SC) adalah pemroses khusus untuk workload jarang. Hal ini bergantung pada hierarki memori tertentu untuk mengelola perpindahan data secara efisien. Compiler XLA berupaya mengukur dan mengalokasikan buffer secara statis berdasarkan batas hardware dan bentuk yang ditentukan pengguna. Error ini menunjukkan kondisi Out of Memory (OOM) selama fase alokasi ini.
Pesan error biasanya menentukan ID ruang memori. Berikut adalah ruang memori umum dan encoding bilangan bulatnya:
| ID Ruang Memori | Nama | Deskripsi |
|---|---|---|
| 0 | Smem | Memori Skalar Lokal. Digunakan untuk register skalar dan alur kontrol. |
| 201 | TileSpmem | Memori Scratchpad Khusus Petak. SRAM lokal yang cepat tersedia untuk petak SC tertentu. |
| 202 | Spmem | Memori Ruang Kerja Bersama. Digunakan untuk mengatur data (input, output, perantara) secara oportunistik untuk menyembunyikan latensi HBM. |
| 203 | HBM | High Bandwidth Memory. Memori bersama yang besar digunakan untuk menyematkan tabel, heap, dan stack. |
| 204 | Tombol Sinkronisasi | Primitif sinkronisasi yang digunakan untuk koordinasi. |
Untuk mempelajari SC dan hierarki memorinya secara mendalam, lihat dokumentasi SparseCore.
Proses debug
Resolusi bergantung pada ruang memori mana yang gagal alokasinya.
Skenario 1. Kegagalan alokasi HBM (ID Ruang Memori: 203)
Error ini terjadi jika alokasi sementara tunggal yang diminta oleh program SparseCore terlalu besar untuk dimuat dalam HBM yang tersedia. Dalam beban kerja penyematan standar dan kolektif yang di-offload SC, partisi per inti yang sangat besar atau spesifikasi sharding yang salah dapat memaksa compiler untuk meminta buffer yang sangat besar.
Tindakan yang disarankan:
- Periksa sharding: Pastikan tabel penyematan dan tensor input/output SC Anda dipartisi/di-shard dengan benar. Jika satu core bertanggung jawab atas terlalu banyak data, alokasi mungkin gagal.
- Menyesuaikan batas: Tinjau
max_ids_per_partitiondanmax_unique_ids_per_partition. Jika nilai ini ditetapkan terlalu tinggi tanpa alasan, compiler akan mencadangkan lebih banyak memori daripada yang diperlukan. Lihat Cara batas diterjemahkan ke dalam tabel.
Skenario 2. Kegagalan memori internal (ID Ruang Memori: 0, 201, 202, 204)
Kegagalan alokasi di Smem, TileSpmem, Spmem, atau Sync Flags biasanya terjadi karena bug atau batasan compiler dalam strategi alokasi, yang menyebabkan compiler gagal memperhitungkan semua persyaratan memori.
Tindakan yang disarankan:
- Isolasi operasi XLA yang gagal: Untuk mengidentifikasi kernel SC HLO atau Mosaic tertentu yang menyebabkan kegagalan, buat representasi compiler perantara:
- Dump SparseCore MLIR: Tetapkan flag
--xla_sc_dump_mlir_to=/path/to/dump. Hal ini akan menghasilkan MLIR program SparseCore, sehingga Anda dapat melihat ukuran alokasi mana yang cocok dengan pesan error. - Dump Mosaic LLO: Untuk kernel kustom, gunakan
--xla_mosaic_dump_to=/path/to/dumpuntuk memeriksa semua program Pengoptimal Tingkat Rendah (LLO) yang dikeluarkan oleh Mosaic.
- Dump SparseCore MLIR: Tetapkan flag
- Kurangi ukuran scratch (pengguna Pallas): Jika kegagalan terjadi dalam kernel Mosaic, tinjau konfigurasi
scratch_shapesAnda. Pastikan permintaanpltpu.SMEMAnda sesuai dengan spesifikasi hardware untuk generasi TPU tertentu. - Nonaktifkan offload kolektif: Jika error muncul dari operasi kolektif yang di-offload SC, coba nonaktifkan fitur offload SC:
--xla_tpu_enable_sparse_core_collective_offload_all_gather=false--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
- Laporkan bug: Jika langkah-langkah di atas tidak menyelesaikan masalah, kemungkinan ada bug compiler. Harap buat laporan bug.