Menyederhanakan Inferensi LLM di Edge dengan TFLite

AGU 13, 2024
Quentin Khan Software Engineer
Linkun Chen Software Engineer

Mengoptimalkan Waktu ke Token Pertama dan Penggunaan Memori Puncak dengan Cache yang Lebih Cerdas untuk XNNPack


XNNPack adalah mesin inferensi CPU TensorFlow Lite default untuk semua model. Ia memberikan peningkatan kecepatan yang sangat signifikan pada platform seluler, desktop, dan Web. Salah satu pengoptimalan yang digunakan dalam XNNPack adalah mengemas ulang bobot statis operator Convolution, Depthwise Convolution, Transposed Convolution, dan Fully Connected ke dalam tata letak internal yang dioptimalkan untuk komputasi inferensi. Selama inferensi, bobot yang dikemas ulang tersebut diakses dalam pola berurutan yang ramah terhadap pipeline prosesor.

Pengurangan latensi inferensi memerlukan biaya: pengemasan ulang pada dasarnya membuat salinan ekstra dari bobot di dalam XNNPack. Upaya sebelumnya telah dilakukan untuk mengurangi biaya tersebut dengan menambahkan in-memorycache ke XNNPack. Cache ini memungkinkan pembagian bobot yang dikemas antara interpreter TFLite independen yang akan menjalankan model yang sama secara independen.

Implementasi TFLite XNNPack delegate telah ditingkatkan untuk mengatasi beberapa kekurangan cache saat ini.


1. Cache berada di memori anonim, yang akan berpindah ke disk jika terjadi tekanan pada memori, sehingga menyebabkan performa yang buruk.

2. Ini memerlukan pengemasan ulang bobot awal setiap kali proses dimulai.

3. Karena pengemasan ulang membaca bobot TFLite asli dan menulis ke buffering baru, ini menyebabkan penggunaan memori puncak yang tinggi selama pengemasan.

4. Ini membutuhkan langkah yang cukup panjang dan manajemen siklus proses yang cermat untuk mengaktifkan caching melalui XNNPack delegate dengan benar.

5. Ini tidak memperbolehkan pembagian bobot di seluruh proses.

TFLite XNNPack delegate architecture
.

Antarmuka Penyedia Cache XNNPack Baru

XNNPack telah diupdate dan menyediakan antarmuka yang memungkinkan Anda mengimplementasikan penyedia cache bobot. Penyedia cache bobot berperilaku sebagai kamus yang akan diisi dan dikueri oleh XNNPack untuk mengakses buffering yang dikemas. Berikut adalah fungsi utamanya.

  • look_up mencari kunci buffering yang telah dikemas dan menampilkan ID unik (atau ID khusus yang disediakan untuk NotFound) yang nantinya dapat digunakan untuk mengambil alamat buffering.

  • reserve_space menyiapkan buffering yang dapat digunakan untuk menyimpan informasi dengan ukuran tertentu. Buffering tersebut perlu di-commit menggunakan look_up_or_insert.

  • look_up_or_insert memeriksa bila buffering yang cocok dengan kunci yang diberikan ada di penyedia cache. Jika tidak, data yang diberikan akan di-commit ke penyedia cache. Fungsi ini juga menampilkan ID yang dapat digunakan untuk mengambil alamat buffering.

  • offset_to_addr menampilkan alamat buffering dari ID yang ditampilkan oleh look_up dan look_up_or_insert.

Interaksi antara XNNPack dan penyedia cache bobot diilustrasikan dalam diagram berikut.

The interactions between XNNPack and the weight cache provider
.

Memuat Cache Dari Disk dengan MMAP dalam TFLite Delegate

TFLite Delegate sekarang menggunakan antarmuka baru dan memiliki penyedia cache bobotnya sendiri. Penyedia ini mampu menyimpan dan memuat bobot yang dikemas secara langsung ke/dari disk. TFLite telah memanfaatkan flatbuffer dan pemetaan memori berbasis file untuk waktu yang lama. Kami mengisi kekosongan di sini dengan memanfaatkan teknik yang sama, untuk mendapatkan keuntungan berikut.


Ini meniadakan overhead pengemasan ulang.

Mempertahankan bobot yang dikemas pada disk melewati proses pengemasan ulang yang mahal setiap kali model dimuat. Ini berarti pengurangan yang signifikan dalam latensi startup dan penggunaan memori puncak. Bahkan untuk pembangunan awal, ini menawarkan penghapusan duplikat data yang dikemas dan selanjutnya meningkatkan performa pengemasan dengan menghindari pengemasan ulang data yang sama.


Ini meningkatkan manajemen memori.

mmap memanfaatkan manajemen memori virtual sistem operasi yang memungkinkannya mengoptimalkan performa dan penggunaan memori sistem secara keseluruhan. Dalam kasus kami, ini sangat menguntungkan untuk akses acak file hanya-baca berukuran besar, seperti bobot konstan operasi neural network misalnya.

Dengan data yang dikemas tersimpan di disk, cache XNNPack tidak lagi bergantung pada memori anonim yang rentan terhadap masalah performa di bawah tekanan memori. Sebaliknya, cache ini memanfaatkan manajemen memori virtual sistem operasi untuk pengoperasian yang lebih lancar.

Karena tidak perlu menyalin data antara sistem file dan memori, mmap secara signifikan mengurangi overhead dan mempercepat waktu akses.

Anda bisa menemukan informasi selengkapnya tentang pemetaan file dan penggunaan memori secara langsung dari halaman utama mmap dan bacaan menarik lainnya.


Ini memungkinkan kolaborasi lintas proses.

Pemuatan file berbasis mmap membuka pintu untuk pembagian bobot yang mudah di antara beberapa proses karena setiap ruang alamat virtual proses dipetakan ke halaman memori fisik yang sama. Hal ini tidak hanya mengurangi jejak memori secara keseluruhan karena beberapa proses berbagi memori yang sama, tetapi juga mempercepat pemuatan model secara keseluruhan.

mmap-based file loading architecture
.

Ini menyederhanakan API yang ditampilkan kepada pengguna.

Alih-alih mengharuskan pengguna untuk mengatur dan mengelola objek cache selama masa pakai aplikasi, mereka cukup menyediakan jalur ke file cache.

std::unique_ptr<tflite::Interpreter> interpreter;
// Setup the options for the XNNPack delegate.
TfLiteXNNPackDelegateOptions xnnpack_options = TfLiteXNNPackDelegateOptionsDefault();
xnnpack_options.weight_cache_file_path = "/tmp/cache_file.xnn_cache";
// Create and apply the XNNPack delegate to a TFLite interpreter.
// Static weights will be packed and written into weights_cache on the first run.
// They will be automatically loaded for all other runs.
TfLiteDelegate* delegate = TfLiteXNNPackDelegateCreate(&xnnpack_options);
interpreter->ModifyGraphWithDelegate(delegate);

Mempertahankan Integritas Cache

Untuk menjamin inferensi yang akurat dan efisien, sangatlah penting untuk membuat cache XNNPack invalid dalam kondisi tertentu:

Evolusi Model: jika bobot atau struktur model Anda berubah, data yang di-cache akan menjadi tidak valid dan harus dibuat invalid. Ini berarti menghapus file di jalur cache yang disediakan.

Upgrade XNNPack: update pada algoritme pengemasan internal XNNPack dapat menyebabkan bobot cache yang tidak kompatibel, sehingga cache harus dikomputasi ulang. Untungnya, XNNPack mampu mendeteksi hal ini dan akan mengganti cache yang ada secara otomatis.

Intinya, setiap modifikasi yang bisa memengaruhi bagaimana bobot dikemas atau digunakan oleh XNNPack akan memicu invalidasi cache.


Tolok ukur

Inisialisasi sesi didominasi oleh pengemasan bobot. Untuk LLM, beberapa subgraf menggunakan ulang bobot yang sama. Membangun cache lebih cepat karena fungsi penghapusan duplikat menghindari pengemasan bobot yang sama beberapa kali. Untuk model yang lebih standar, seperti stable diffusion, tidak ada penghapusan duplikat dan waktu inisialisasi yang sedikit lebih tinggi disebabkan oleh penyimpanan cache ke disk. Memuat ulang cache (berjalan dari proses ke-2) membuat inisialisasi menjadi lebih singkat dari waktu sebelumnya dalam semua kasus.

Peningkatan inisialisasi sesi secara otomatis memengaruhi waktu ke token pertama untuk LLM, kira-kira setengahnya dalam tolok ukur.

Peningkatan memori yang dihasilkan oleh implementasi cache juga bisa dilihat. Resident Set Size puncak diturunkan untuk LLM berkat penghapusan duplikat. Untuk model lain yang tidak mendapatkan keuntungan dari penghapusan duplikat, tidak ada perubahan. Memuat ulang cache akan menurunkan RSS puncak lebih banyak karena model asli TFLite tidak dibaca lagi sehingga tidak akan ditarik ke dalam memori.


Gemma 2B di Pixel 8 Pro

Benchmarks - Gemma 2B on a Pixel 8 Pro
.

Phi2 di Pixel 8 Pro

Benchmarks - Phi2 on a Pixel 8 Pro
.

Stable Diffusion di Pixel 8 Pro

Stable Diffusion on a Pixel 8 Pro
.

Karya Selanjutnya

Saat ini cache masih terikat dengan penggunaan sistem file. Kami ingin memanfaatkan mekanisme penghapusan duplikat data secara independen untuk kasus penggunaan tertentu bila pengguna tidak ingin menukar memori yang dialokasikan secara tradisional dengan pemetaan berbasis file. mmap memperbolehkan pembuatan pemetaan anonim yang memungkinkan penggunaan ulang sebagian besar implementasi.