Implementasi Mamba dalam Satu Berkas PyTorch

 (github.com/johnma2006)
1 poin oleh GN⁺ 2023-12-21 | 1 komentar | Bagikan ke WhatsApp
  • mamba-minimal adalah proyek yang mengimplementasikan Mamba secara sederhana dan minimal dalam satu berkas PyTorch
  • Bertujuan menghasilkan keluaran numerik yang sama dengan implementasi resmi pada forward/backward pass
  • Kodenya disederhanakan dan disusun dalam bentuk berkomentar agar mudah dibaca
  • Tidak menyertakan optimasi inti dari implementasi resmi, sehingga tidak menawarkan kecepatan, dan juga tidak mencakup inisialisasi parameter yang tepat
  • Demo menjalankan contoh penyelesaian prompt menggunakan state-spaces/mamba-370m dan tokenizer EleutherAI/gpt-neox-20b

Gambaran proyek

  • mamba-minimal adalah implementasi minimal sederhana dari Mamba dalam satu berkas PyTorch
  • Tujuannya adalah menunjukkan perilaku yang sama seperti implementasi resmi dengan kode yang lebih mudah dibaca
  • Fitur utama:
    • Keluaran numerik yang setara dengan implementasi resmi pada forward pass dan backward pass
    • Kode yang disederhanakan
    • Implementasi yang mudah dibaca dan diberi komentar

Yang tidak disertakan

  • Kecepatan bukanlah tujuannya
    • Implementasi resmi sangat dioptimalkan
    • Optimasi tersebut merupakan bagian dari kontribusi utama makalah Mamba
    • Implementasi ini mempertahankan sebagian besar kode tetap sederhana demi keterbacaan
  • Inisialisasi parameter yang tepat tidak disertakan
    • Disebutkan sebagai hal yang bisa ditambahkan tanpa mengorbankan keterbacaan

Contoh penggunaan demo

  • Contoh penyelesaian prompt dapat dilihat di demo.ipynb
  • Contoh menggunakan model.Mamba dan AutoTokenizer dari Hugging Face transformers
  • Model dan tokenizer yang digunakan:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • Contoh prompt adalah Mamba is the, dan hasil generasinya mencakup kalimat yang menjelaskan Mamba sebagai ular berbisa

Referensi

Bacaan terkait

1 komentar

 
GN⁺ 2023-12-21
Komentar Hacker News

  • Dulu saya membuat library bersama seorang rekan yang memisahkan sebagian besar kode model bersama, dan dengan ini banyak model bisa diimplementasikan dalam sekitar 100 baris jika tidak menghitung Python import dan komentar
    BERT: https://github.com/explosion/curated-transformers/blob/main/...
    Llama 1/2: https://github.com/explosion/curated-transformers/blob/main/...
    MPT: https://github.com/explosion/curated-transformers/blob/main/...
    Juga mendukung fitur seperti TorchScript JIT dan PyTorch flash attention

    • Saya pasti akan melihat library ini. Penasaran apakah Anda juga sudah melihat xformers
      xformers menangani masalah yang mirip, tetapi lebih berfokus menyediakan modul Transformer berperforma tinggi dengan Triton. Namun, tidak mudah mengambil dan memakai hanya komponen tertentu dari library itu, dan saya terus menemui runtime error jadi untuk sementara saya tunda. Saya sedang membuat sesuatu berbasis arsitektur BERT, jadi akan saya jadikan referensi
    • Saya terkesan dengan library ini. Saya tidak terlalu menyukai implementasi Hugging Face, tetapi ini tampak seperti API yang indah dengan tingkat abstraksi yang pas
      Saya berencana mencobanya di proyek berikutnya

  • Kode Mamba aslinya sulit langsung dipahami karena banyak optimasi kecepatan dan elemen lainnya, tetapi implementasi ini sepertinya akan membantu untuk belajar
    Saat melakukan inferensi token satu per satu, semuanya menjadi jauh lebih sederhana. Saya juga punya implementasi inferensi Mamba buatan sendiri: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran, ya. Saya penasaran kenapa Anda memakai Fortran
      Saya tahu itu adalah fondasi kode komputasi ilmiah yang sudah lama teruji dan sering dibungkus oleh library seperti PyTorch atau Numpy, tetapi sekarang bukan bahasa yang populer. Saya penasaran apa alasan Anda memilihnya

  • Ada bagian yang saya harap bisa menjelaskan Mamba dengan cara yang bisa dipahami bahkan oleh orang yang bukan peneliti machine learning

    1. Apa wawasan umum yang dimiliki model ruang-keadaan di luar Transformer
    2. Inovasi bertahap apa yang membuat Mamba lebih sukses atau lebih menarik dibanding pendahulunya seperti S4, H3, dan Monarch
    3. Selain skalabilitas subkuadratik terhadap panjang konteks, apa lagi maknanya. Misalnya, jika kita tidak tertarik pada panjang konteks di atas 100k token, saya penasaran apakah pada model dan dataset dengan ukuran serupa Mamba berpotensi lebih efisien secara komputasi saat pelatihan
    • Kecerdasan saya jauh di bawah para penulis makalah ini, tapi saya tetap berusaha memahaminya. Saya belajar ilmu komputer dan punya dasar teori kontrol serta intuisi sistem waktu-diskret tingkat sarjana, tetapi tampaknya untuk benar-benar memahami makalah ini saya perlu mempelajari model ruang-keadaan jauh lebih dalam
      Wawasan inti Mamba terletak pada penyelesaian masalah lama pada model ruang-keadaan. Model ruang-keadaan bagus untuk mengompresi konteks masukan, tetapi dalam proses mengompresi masukan menjadi keadaan tersembunyi, informasi yang dibutuhkan untuk memanfaatkan konteks secara efektif seperti pada Transformer menjadi terhapus
      Solusinya adalah membuat apa yang dalam makalah disebut mekanisme seleksi. Mekanisme ini bergantung pada masukan, sehingga setiap kali masukan berubah, model dapat menyesuaikan keluaran pada tiap langkah. Untuk itu, beberapa variabel ruang-keadaan dibuat bergantung pada masukan alih-alih invarian terhadap masukan, dan lapisan linear serta sejenisnya ditambahkan untuk memproyeksikan masukan pada tiap waktu ke variabel ruang-keadaan
      Namun, membuat variabel ruang-keadaan bergantung pada masukan menimbulkan overhead komputasi. Ini diatasi dengan algoritme sadar-perangkat-keras yang memanfaatkan struktur memori GPU modern semaksimal mungkin, sehingga perpindahan data masuk dan keluar HBM dihindari sebisa mungkin
      Tri Dao adalah orang yang membuat Flash Attention, dan itu juga merupakan cara untuk menggunakan perangkat keras secara lebih efisien pada Transformer. Bidang ini memang benar-benar keahliannya
    • Attention tumbuh secara kuadratik terhadap panjang konteks, sedangkan jaringan saraf berulang dengan gating (LSTM, GRU, dan sebagainya) bersifat linear, dan arsitektur baru ini juga linear. Jaringan saraf berulang generasi awal memakai gating untuk menghindari gradien yang meledak, tetapi pendekatan baru menggunakan teori sistem dinamis untuk menjamin stabilitas, sehingga gating tidak perlu menyelesaikan dua masalah sekaligus dan bisa fokus pada memori
      Mamba dan Based, yang muncul tepat sebelum NeurIPS 2023, memasukkan multi-query associative recall (MQAR) serta ketergantungan data pada gating/seleksi yang terinspirasi dari multi-head Attention. Keduanya adalah unsur kunci yang hilang pada Hyena dan arsitektur ruang-keadaan sebelumnya, dan membuat model baru ini menjadi sebaik Attention pada tugas associative recall, serta mungkin sedikit lebih baik daripada Attention pada tugas lain yang bukan sekadar lookup
      Tentu saja, detail besar pada Mamba adalah implementasi CUDA yang efisien. Tanpa itu, arti penting arsitektur ini bisa berkurang pada tugas-tugas yang memang sudah cocok untuk Transformer
      Bahkan jika Anda tidak terlalu mengkhawatirkan panjang konteks, banyak wilayah baru tetap terbuka. Analisis sekuens DNA adalah tugas linear dengan ketergantungan panjang, dan kita juga bisa memikirkan citra, video, serta informasi berdimensi tinggi sebagai aliran token. Mirip seperti memindai piksel pada monitor CRT lama
      Salah satu impian awal AI adalah lintasan pembelajaran tunggal dari agen yang terus berinteraksi dengan lingkungan dan berevolusi terus-menerus, dan model dengan panjang konteks tak terbatas seperti ini mungkin bisa mempermudah terwujudnya impian tersebut
      Namun untuk saat ini, aplikasi hilir pada tugas nyata yang penting untuk model semacam ini umumnya masih kurang tervalidasi dan kurang dituning dibanding aplikasi berbasis Attention yang sudah matang. Analogi dengan jaringan saraf berulang lama cukup membantu sampai batas tertentu, tetapi selama 5 tahun terakhir orang terlalu terspesialisasi pada Attention dan Transformer, sehingga inersia ke arah Transformer sangat kuat
    • Saya juga ingin tahu apakah Mamba bisa dilatih dengan efisiensi komputasi yang lebih baik pada model dan dataset berukuran serupa
      Dalam makalah aslinya dijelaskan bahwa setelah parameter ditransformasikan, model dapat dihitung dengan dua cara: relasi rekurensi linear atau konvolusi global. Umumnya, saat pelatihan ketika seluruh sekuens masukan bisa dilihat lebih dulu, digunakan mode konvolusi yang mudah diparalelkan, sedangkan pada inferensi autoregresif yang melihat masukan satu langkah waktu demi satu langkah waktu, model beralih ke mode rekuren yang efisien
      Jadi pelatihan dapat diparalelkan, mirip mode forward paralel pada RetNet. Inferensi dasarnya dijalankan dalam mode rekuren untuk memperoleh konteks sepanjang mungkin, dan karena tidak ada chunking, sulit menilai seberapa besar RAM dan VRAM yang akan dipakai selama inferensi
    • Sepertinya video ini sangat cocok dengan yang Anda cari
      Video ini menjelaskan makalahnya sambil juga memberi banyak konteks tentang posisinya dalam gambaran besar. Alur penjelasannya cukup menarik untuk diikuti
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • Sepengetahuan saya, Mamba pada dasarnya adalah kelanjutan dari riset model ruang-keadaan yang bisa disebut sebagai konvolusi panjang
      Alih-alih Attention kuadratik yang menghitung seberapa besar setiap token memperhatikan semua token lain, pendekatannya adalah entah bagaimana menghitung kernel konvolusi panjang yang panjangnya sama dengan masukan, lalu menerapkan conv1d
      Dari pemahaman saya yang terbatas, ini sedikit berkaitan dengan menerapkan FFT, melakukan perkalian matriks, lalu kembali dengan IFFT. Saya tahu itu bekerja, tetapi lambat. Ada berbagai cara menghitung FFT, dan salah satunya adalah matriks kupu-kupu. Mungkin itu hanya pendekatan aproksimasi, tetapi tampaknya cukup bagus dan sangat cepat serta efisien pada perangkat keras saat ini
      Kompleksitas kuadratik terdengar buruk, tetapi dalam praktiknya algoritme subkuadratik sering kali lebih lambat karena keterbatasan perangkat keras. Jadi, meskipun ekspektasi terhadap model ruang-keadaan besar, masih sulit mengatakan bahwa Llama sudah selesai. Kita juga belum tahu apakah Mamba tetap bekerja baik saat diskalakan besar, dan untuk mengetahuinya perlu benar-benar menghabiskan jutaan dolar untuk pelatihan. Meski begitu, saya tetap optimistis
      Model lain yang menarik dalam keluarga subkuadratik adalah RWKV. Layak dilihat, tetapi sepertinya itu sudah pernah dibahas di podcast
      Saya belajar sendiri dan hanya pernah menelusuri makalahnya secara sepintas dulu, jadi mungkin banyak yang salah. Selain itu, Attention biasanya punya cache KV yang sangat membantu performa, dan menurut saya Mamba tidak bisa melakukan hal itu

  • Saya tertawa pada kalimat “Mamba adalah ular berbisa terpanjang di dunia dengan panjang yang diperkirakan melebihi 150 m”
    Tetap saja, ini benar-benar luar biasa, dan karena merujuk ke makalah arXiv, orang seperti saya yang mengonsumsi tulisan seperti ini alih-alih menafsirkan makalahnya langsung jadi bisa sedikit mengintip bagian dalamnya

    • Nama Mamba bagus. Karena itu adalah [S]elective [S]tructured [S]tate [S]pace [S]equence models, jadinya sSSSS, dan terdengar seperti suara ular
    • Saya kira ular berbisa terpanjang itu king cobra. Cari cepat di Google juga menunjukkan begitu
      Akan lucu juga kalau nanti perlu menerbitkan koreksi untuk kalimat itu

  • Saya menduga inti algoritmenya adalah parallel prefix scan. Rasanya itulah poin utama Mamba
    for i in range(l):
    x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
    y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')
    ys.append(y)

  • Mungkin ini pertanyaan bodoh, tapi saya penasaran seberapa sulit melatih model Mamba yang diunggah ke Hugging Face
    Model terbesarnya tampak 2.8B; kalau dilatih dengan dataset seperti The Pile, kira-kira butuh berapa GPU dan berapa lama?

    • Pertanyaan bagus yang juga ingin saya ketahui jawabannya. Jawabannya tampaknya jauh lebih cepat daripada Transformer dengan ukuran yang sama, dan hasil akhirnya juga kemungkinan mendapat skor lebih baik daripada Transformer di hampir semua benchmark
      Inferensinya juga tampaknya 3–5x lebih cepat sambil hanya memakai setengah RAM

  • Saya sempat mencoba membedah versi CUDA resminya, tapi setelah percobaan pertama gagal akhirnya tidak jadi saya sentuh lagi, dan implementasi ini terlihat jauh lebih baik

  • Satu lagi implementasi PyTorch satu file, benar-benar luar biasa. Saya berharap hlb-CIFAR10 dan proyek terkait yang saya lakukan sebelumnya, serta pengaruh pendahulu seperti minGPT dan DawnBench, sedikit banyak membantu mendorong format satu file yang sederhana
    Pekerjaan seperti ini penting untuk riset machine learning yang efisien, dan mungkin merupakan salah satu hal terpenting yang bisa dilakukan untuk bidang ini saat ini
    Riset bergerak maju dengan kecepatan inovasi, dan inovasi makin cepat sebanding dengan kebalikan dari waktu eksekusi eksperimen, yang jelas terkait dengan kompleksitas Kolmogorov dari kode untuk keperluan riset atau hacking sederhana
    Sulit untuk melebih-lebihkan betapa pentingnya alat seperti ini bagi riset, dan seberapa besar secara pribadi ini mempercepat proses penemuan pengetahuan bagi saya. Kemampuan untuk dengan cepat membuat sketsa ide dalam hitungan menit dan segera mendapatkan hasil dengan rasio signal-to-noise yang tinggi telah menjadi unsur penting dalam kemajuan riset
    Saya melihat distilasi pengetahuan dan MDL(https://en.wikipedia.org/wiki/Minimum_description_length) sebagai hal yang sangat penting untuk membalikkan ornamen yang tidak perlu, sampah, dan persaingan topik bernilai rendah yang terlalu padat demi “tidak ketinggalan”, yang tampaknya didorong oleh proses submit/review makalah saat ini
    Belakangan ini, untuk menghindari masalah itu dan bergerak ke solusi penskalaan yang sedikit lebih baik, saya mulai merilis kode sebagai “code sketch”, yaitu gist pendek satu file yang mandiri. Ini mengurangi waktu pengembangan, dan memungkinkan saya langsung menyampaikan kode kerja yang kasar dan belum dipoles yang memuat konsepnya kepada orang lain. Sejauh ini tampaknya cukup berhasil dan saya ingin terus melakukannya
    Saya ingin melihat lebih banyak kode seperti ini. Jika Anda peneliti yang melatih data dalam skala besar, Anda juga harus efisien terhadap data dalam cara menyebarkan informasi

    • Tahun 2023 adalah tahun yang menarik hanya dengan melihat riset AI berkembang dengan kecepatan yang tidak masuk akal. Elemen-elemen dasar seperti ArXiV, PyTorch, GitHub, Hugging Face, dan kode Python open source yang ringkas secara dramatis mempercepat perkembangan bidang baru ini
      Mungkin belum pernah ada manusia yang mengembangkan sesuatu dengan kompleksitas sebesar ini secepat ini
      Satu-satunya tempat lain yang tampak memiliki kecepatan serupa mungkin SpaceX, yang tahun ini juga meluncurkan dua roket mutakhir. Penasaran apa yang akan muncul pada 2024
    • Mungkin ada peluang peningkatan performa kecil. Di sini x_proj tidak memiliki bias, jadi tampaknya bobot x_proj dan dt_proj bisa digabungkan
      Jika ada kebutuhan penyesuaian bobot, mungkin itu bisa dilakukan begitu saja saat runtime, dan kernel tunggal plus bias pada akhirnya mungkin lebih cepat. Tidak yakin juga

  • Saya penasaran apakah ada diskusi untuk makalah aslinya. Mungkin saya melewatkannya, tapi ini cukup menarik
    Saya kurang paham bagian yang mengatakan “karena kurangnya implementasi efisien menimbulkan kehabisan memori atau kebutuhan komputasi yang tidak realistis, hasil penuh pada panjang konteks 8k untuk baseline RWKV dan RetNet, model rekuren kuat sebelumnya yang juga dapat ditafsirkan sebagai SSM, tidak disertakan”
    RetNet tidak banyak memakai memori, dan jika menggunakan implementasi forward berbasis chunk, penggunaan VRAM dibatasi oleh ukuran chunk. Bagian inilah yang penting saat menguji panjang konteks
    Saya penasaran apakah ada yang sudah menguji model Mamba asli. Dibandingkan dengan RetNet dalam mode forward paralel, kira-kira bagaimana kecepatan latihnya?

  • Implementasi yang merangkum hal kompleks hanya menyisakan intinya selalu bagus