Kode error: E2001

Kategori: Waktu Kompilasi: Jenis Data RHS yang Tidak Didukung di Hardware

Error ini terjadi saat jenis data yang digunakan untuk operand Sisi Kanan (RHS) dalam perkalian matriks (misalnya, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul, atau operator @) tidak didukung secara native oleh generasi TPU tertentu yang digunakan.

Contoh pesan error:

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

Backend XLA: TPU

Ringkasan

Matrix Multiply Unit (MXU) TPU secara native mendukung operasi Float32 di semua generasi hardware.

Namun, dukungan native untuk BFloat16 dan jenis data terkuantisasi lainnya (misalnya, Int4, Int8, atau Float8) bervariasi menurut generasi hardware. Error ini dipicu saat kernel Anda mencoba memetakan perkalian matriks ke MXU menggunakan jenis data yang tidak memiliki sirkuit fisik untuk dieksekusi oleh generasi TPU tertentu Anda.

Error ini biasanya menunjukkan bahwa proses Kanonisasi compiler— yang mencoba mengonversi jenis yang tidak didukung menjadi jenis yang didukung secara otomatis (misalnya, melalui emulasi software)—tidak dapat menemukan aturan konversi yang valid atau dicegah melakukannya karena Mode Kompatibilitas dinonaktifkan.

Proses debug

Untuk mengatasi error ini, Anda harus menyelaraskan jenis data dengan kemampuan hardware Anda. Anda memiliki opsi berikut:

1. Mentransmisikan ke jenis native

Perbaikan yang paling andal adalah dengan melakukan casting operand secara manual ke jenis data yang didukung hardware (seperti Float32 atau BFloat16 di TPU v4+) di dalam kernel sebelum operasi matmul.

  • Alasan: Float32 adalah jenis data universal yang didukung secara native oleh MXU di semua generasi TPU.
  • Trade-off: Hal ini disertai dengan biaya VPU (Vector Processing Unit) - siklus yang diperlukan untuk melakukan cast, tetapi menjamin kernel Anda akan berjalan di hardware saat ini.

2. Memeriksa Mode Kompatibilitas

Biasanya, compiler dapat menangani masalah ketidakcocokan jenis ini secara otomatis dalam Mode Kompatibilitas yang diaktifkan secara default. Periksa kembali konfigurasi XLA untuk memastikan --xla_mosaic_compat_mode tidak disetel ke salah (false).

Hal ini berfungsi sebagai "polyfill", yang menyuntikkan urutan emulasi software untuk operasi yang tidak didukung secara native oleh hardware Anda.

Yang dapat dilakukan Mode Kompatibilitas:

  • MatMul presisi campuran: Memungkinkan pencampuran operan Integer dengan akumulator Float dengan otomatis menyisipkan operasi transmisi (misalnya, memperluas bilangan bulat ke Float32 sebelum matmul).
  • Emulasi presisi rendah: Pada generasi hardware tertentu, mengemulasi jenis yang tidak didukung seperti floating point 4-bit (4E2M1FN) atau floating point 8-bit (8E4M3FN) dengan memperluasnya ke jenis yang didukung seperti BFloat16 atau Float32 sebelum dieksekusi.

Perhatikan bahwa mode ini memprioritaskan kompatibilitas daripada performa puncak karena emulasi memerlukan petunjuk tambahan untuk mengonversi format data sebelum MXU dapat mengoperasikannya.

3. Mengupgrade hardware atau meminta dukungan

Jika algoritma Anda benar-benar memerlukan performa native untuk jenis seperti Int4 atau Float8 tanpa overhead casting atau emulasi, Anda harus menjalankan algoritma di generasi TPU yang lebih baru dengan dukungan native.

Permintaan fitur: Jika Anda yakin hardware Anda mendukung operasi ini, atau jika compiler tidak memiliki jalur emulasi yang valid meskipun dalam Mode Kompatibilitas, ajukan permintaan fitur. Kami biasanya menjamin bahwa operasi kompatibel dengan versi yang lebih baru. Jadi, jika kernel Anda berjalan di generasi TPU, kernel tersebut akan berjalan di semua generasi mendatang, tetapi tidak dijamin memiliki emulasi untuk generasi lama (yang beberapa di antaranya akan sangat mahal).