PyTorch’ta tek dosyada Mamba implementasyonu

 (github.com/johnma2006)
1 puan yazan GN⁺ 2023-12-21 | 1 yorum | WhatsApp'ta paylaş
  • mamba-minimal, Mamba’nın PyTorch’ta tek dosyalık basit ve minimal bir implementasyonunu sunan bir projedir
  • Resmî implementasyonla forward/backward pass sırasında aynı sayısal çıktıları üretmeyi hedefler
  • Kod sadeleştirilmiştir ve yorumlar eklenmiş, okunması kolay bir yapı ile hazırlanmıştır
  • Resmî implementasyondaki temel optimizasyonları içermediği için hız sağlamaz ve uygun parametre başlatmayı da içermez
  • Demo, state-spaces/mamba-370m ve EleutherAI/gpt-neox-20b tokenizer’ını kullanarak istem tamamlama örneği çalıştırır

Proje özeti

  • mamba-minimal, Mamba’nın PyTorch’ta tek dosyalık basit bir minimal implementasyonudur
  • Amaç, resmî implementasyonla aynı davranışı daha okunabilir bir kodla göstermektir
  • Başlıca özellikler:
    • Resmî implementasyonla forward pass ve backward pass için eşdeğer sayısal çıktı
    • Sadeleştirilmiş kod
    • Okunması kolay ve yorum eklenmiş implementasyon

İçermedikleri

  • Hız hedeflenmez
    • Resmî implementasyon yoğun biçimde optimize edilmiştir
    • Bu optimizasyonlar Mamba makalesinin temel katkılarının parçasıdır
    • Bu implementasyon, okunabilirlik için çoğu kısmı bilinçli olarak sade tutar
  • Uygun parametre başlatma içermez
    • Okunabilirlikten ödün vermeden eklenebilecek bir unsur olarak belirtilir

Demo kullanım örneği

  • demo.ipynb içinde istem tamamlama örneği görülebilir
  • Örnek, model.Mamba ve Hugging Face transformers içindeki AutoTokenizer kullanır
  • Kullanılan model ve tokenizer:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • Örnek istem Mamba is the şeklindedir ve üretim çıktısında Mamba’yı engerek olarak tanımlayan bir cümle yer alır

Referanslar

İlgili okumalar

1 yorum

 
GN⁺ 2023-12-21
Hacker News yorumları

  • Bir süre önce bir iş arkadaşımla, paylaşılan model kodunun çoğunu ayıran bir kütüphane oluşturmuştuk; bunu kullanınca birçok modeli Python import'ları ve yorumlar hariç yaklaşık 100 satırda uygulamak mümkün
    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/...
    TorchScript JIT ve PyTorch flash attention gibi özellikleri de destekliyor

    • Bu kütüphaneye mutlaka göz atacağım. xformers'a da bakıp bakmadığınızı merak ediyorum
      xformers da benzer bir sorunu ele alıyor, ancak Triton kullanarak yüksek performanslı Transformer modülleri sunmaya daha fazla odaklanıyor. Yine de kütüphanenin belirli bileşenlerini alıp kullanmak pek kolay değildi ve sürekli çalışma zamanı hataları aldığım için şimdilik erteledim. BERT mimarisi temelinde bir şeyler geliştirdiğim için referans olarak inceleyeceğim
    • Bu kütüphaneye hayran kaldım. Hugging Face implementasyonlarını pek sevmiyordum; bu ise soyutlama seviyesi tam yerinde olan güzel bir API gibi görünüyor
      Bir sonraki projemde denemeyi düşünüyorum

  • Orijinal Mamba kodu, hız optimizasyonları ve başka unsurlar çok olduğu için doğrudan anlaması zor; bu implementasyon öğrenmeye yardımcı olacak gibi
    Token'ları tek tek çıkarım sırasında işlerken her şey çok daha basitleşiyor. Kendi yaptığım bir Mamba çıkarım implementasyonum da var: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran mı? Neden Fortran kullandığınızı merak ediyorum
      Uzun zamandır doğrulanmış bilimsel hesaplama kodlarının temeli olduğunu ve PyTorch ya da Numpy gibi kütüphanelerle sarılıp kullanıldığını biliyorum, ama günümüzde popüler bir dil değil. Neden seçtiğinizi merak ediyorum

  • Mamba hakkında, makine öğrenimi araştırmacısı olmayan birinin de anlayabileceği şekilde açıklanmasını istediğim noktalar var

    1. Transformer ötesindeki durum-uzayı modellerinin genel içgörüsü nedir?
    2. Mamba’yı S4, H3, Monarch gibi önceki çalışmalardan daha başarılı ya da ilginç kılan kademeli yenilik nedir?
    3. Bağlam uzunluğunda karesel altı ölçeklenebilirlik dışında ne anlamı var? Örneğin 100k token’dan uzun bağlam uzunluklarıyla ilgilenmiyorsam, benzer büyüklükteki model ve veri kümelerinde Mamba’nın eğitim hesaplama verimliliğinin daha iyi olma ihtimali var mı merak ediyorum
    • Zekâm makalenin yazarlarının çok gerisinde ama yine de anlamaya çalıştım. Bilgisayar bilimi okudum; temel kontrol teorisi ve lisans düzeyinde ayrık zamanlı sistem sezgim var, ancak bu makaleyi doğru düzgün anlamak için durum-uzayı modellerini çok daha fazla çalışmam gerekecek gibi görünüyor
      Mamba’nın temel içgörüsü, durum-uzayı modellerinin eski bir sorununu çözmesinde yatıyor. Durum-uzayı modelleri girdi bağlamını sıkıştırmada iyi, ancak girdiyi gizli duruma sıkıştırma sürecinde Transformer gibi bağlamı etkili kullanmak için gereken bilgiler siliniyor
      Çözüm, makalede seçim mekanizması denen şeyi oluşturmak. Bu mekanizma girdiye bağımlı olduğu için, girdi her değiştiğinde model her adımdaki çıktısını ayarlayabiliyor. Bunun için bazı durum-uzayı değişkenlerini girdiden bağımsız değil girdiye bağımlı hâle getiriyor ve her zaman adımındaki girdiyi durum-uzayı değişkenlerine projekte edecek doğrusal katmanlar vb. ekliyor
      Ancak durum-uzayı değişkenlerini girdiye bağımlı yapmak hesaplama ek yükü doğuruyor. Bunu, modern GPU bellek yapısından azami yararlanan donanım farkındalıklı bir algoritmayla çözüp veriyi HBM’in içine ve dışına taşımayı mümkün olduğunca önlüyorlar
      Tri Dao, Flash Attention’ın yaratıcısı; bu da Transformer’da donanımı daha verimli kullanmanın bir yoluydu. Bu alan gerçekten onun uzmanlığı
    • Attention, bağlam uzunluğuna göre ikinci dereceden büyür; kapılamalı yinelemeli sinir ağları (LSTM, GRU vb.) doğrusaldır ve bu yeni mimariler de doğrusaldır. İlk yinelemeli sinir ağları patlayan gradyanlardan kaçınmak için kapılama kullanıyordu; yeni yaklaşımlar ise kararlılığı garanti eden dinamik sistemler teorisini kullanarak kapılamanın iki sorunu aynı anda çözmek yerine belleğe odaklanmasını sağlıyor
      NeurIPS 2023’ten hemen önce çıkan Mamba ve Based, çoklu sorgu ilişkisel hatırlama (MQAR) ile çok başlı Attention’dan esinlenen kapılama/seçimin veriye bağımlılığını içeriyordu. Hyena ve önceki durum-uzayı mimarilerinde eksik olan temel unsurlar bunlardı; yeni modeller ilişkisel hatırlama görevlerinde Attention kadar iyi hâle geliyor ve arama dışındaki görevlerde belki Attention’dan biraz daha iyi olma ihtimali de gösteriyor
      Elbette Mamba’nın büyük ayrıntısı verimli CUDA uygulaması. O olmazsa Transformer’ın zaten uygun olduğu işlerde bu mimarinin anlamı azalabilir
      Bağlam uzunluğunu çok dert etmeseniz bile birçok yeni alan açılıyor. DNA dizi analizi uzun bağımlılıkları olan doğrusal bir iş; görüntü, video ve yüksek boyutlu bilgileri token akışı olarak görme biçimleri de düşünülebilir. Eski CRT monitörler gibi pikselleri taramak gibi
      Yapay zekanın ilk hayallerinden biri, çevresiyle sürekli etkileşen bir ajanın tek bir öğrenme yörüngesinin durmadan evrilmesiydi; bu tür sonsuz bağlam uzunluğu modelleri bu hayali kolaylaştırabilir
      Ancak şu an için bu modellerin önemli gerçek işlerdeki downstream uygulamaları, Attention tabanlı olgun uygulamalara kıyasla genel olarak daha az doğrulanmış ve ayarlanmış durumda. Eski yinelemeli sinir ağlarıyla benzetme bir ölçüde yardımcı oluyor, fakat son 5 yılda insanlar Attention ve Transformer’a fazlasıyla özelleştiği için Transformer tarafında büyük bir atalet var
    • Benzer büyüklükte model ve veri kümelerinde Mamba’nın daha hesaplama verimli eğitilip eğitilemeyeceğini ben de bilmek isterim
      Orijinal makalede, parametreler dönüştürüldükten sonra modelin doğrusal özyineleme ya da küresel evrişim olmak üzere iki şekilde hesaplanabileceği açıklanıyor. Genel olarak tüm girdi dizisini önceden görebildiğiniz eğitimde paralelleştirmesi kolay olan evrişim modu kullanılıyor; girdiyi tek tek zaman adımlarıyla gören otoregresif çıkarımda ise verimli yinelemeli moda geçiliyor
      Bu yüzden eğitim, RetNet’in paralel ileri geçiş modu gibi paralelleştirilebilir. Temel çıkarım, mümkün olduğunca uzun bağlam elde etmek için yinelemeli modda yapılıyor; parçalama olmadığı için çıkarım sırasında RAM ve VRAM’i ne kadar tüketeceğini değerlendirmek zor
    • Bu video aradığın şeye tam uyacak gibi görünüyor
      Makaleyi açıklarken büyük resimde nereye oturduğuna dair de bolca bağlam veriyor. Anlatımı dinlemek oldukça ilginç
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • Bildiğim kadarıyla Mamba, temelde uzun evrişim diye adlandırılabilecek durum-uzayı modeli araştırmasının devamı
      Her token’ın diğer tüm token’lara ne kadar dikkat edeceğini hesaplayan ikinci dereceden Attention yerine, bir şekilde girdiyle aynı uzunlukta uzun bir evrişim çekirdeği hesaplayıp conv1d uyguluyor
      Sınırlı anladığım kadarıyla bunun FFT uygulamak, matris çarpımı yapmak ve sonra tekrar IFFT ile geri dönmekle biraz ilgisi var. Çalıştığını biliyorum ama yavaş. FFT hesaplamanın birçok yolu var ve bunlardan biri kelebek matrisi. Muhtemelen yalnızca bir yaklaşım ama yeterince iyi ve mevcut donanımda çok hızlı ve verimli görünüyor
      İkinci dereceden karmaşıklık kulağa kötü geliyor, ancak pratikte donanım kısıtları yüzünden karesel altı algoritmalar çoğu zaman daha yavaş olabiliyor. Bu yüzden durum-uzayı modellerine yönelik beklenti büyük olsa da Llama’nın bittiğini söylemek kolay değil. Mamba’nın ölçeği büyütüldüğünde de iyi çalışıp çalışmayacağını henüz bilmiyoruz; bunu öğrenmek için gerçekten eğitime milyonlarca dolar harcamak gerekir. Yine de iyimserim
      Karesel altı ailesindeki bir başka ilginç model RWKV. Bakmaya değer ama muhtemelen podcast’te daha önce ele alınmıştır
      Kendi kendime öğrendim ve makaleye de eskiden üstünkörü göz atmıştım; bu yüzden çok yanılıyor olabilirim. Ayrıca Attention’da genellikle KV cache bulunur ve bu performansa büyük katkı sağlar; Mamba’da bunun yapılamayacağını düşünüyorum

  • “Mamba, tahmini uzunluğu 150 metreyi aşan dünyanın en uzun zehirli yılanıdır” cümlesine güldüm.
    Yine de gerçekten çok iyi; arXiv makalesine referans verdiği için, makaleyi doğrudan yorumlamak yerine böyle yazıları tüketen benim gibi insanlar da içeriye biraz göz atabiliyor.

    • Mamba adı güzel. [S]elective [S]tructured [S]tate [S]pace [S]equence models olduğu için sSSSS oluyor ve yılan sesi gibi duyuluyor.
    • En uzun zehirli yılanın kral kobra olduğunu sanıyordum. Basit bir Google araması da öyle gösterdi.
      İleride o cümle için bir düzeltme yayımlanırsa eğlenceli olur.

  • Algoritmanın özünün paralel önek taraması olacağını tahmin etmiştim. Mamba’nın asıl meselesi bu değil mi diye düşünüyorum.
    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)

  • Aptalca bir soru olabilir ama Hugging Face’teki Mamba modelini eğitmenin zorluğunu merak ediyorum.
    En büyük model 2.8B gibi görünüyor; The Pile gibi bir veri kümesiyle eğitmek için kaç GPU gerekir ve ne kadar sürer?

    • Benim de bilmek istediğim harika bir soru. Yanıt, aynı boyuttaki Transformer’dan oldukça hızlı olduğu yönünde gibi; nihai sonuçlar da neredeyse tüm benchmark’larda Transformer’dan daha iyi puan alacak gibi görünüyor.
      Çıkarımın da RAM’in yalnızca yarısını kullanarak 3-5 kat daha hızlı yapılacağı anlaşılıyor.

  • Resmî CUDA sürümünü çözmeye çalışmış, ilk denemem başarısız olduktan sonra sonunda elleyememiştim; bu uygulama çok daha iyi görünüyor.

  • Bir başka tek dosyalık PyTorch uygulaması, gerçekten harika. Daha önceki hlb-CIFAR10 ve ilgili projelerin, ayrıca minGPT ya da DawnBench gibi öncül etkilerin basit tek dosya biçimini az da olsa ileri taşımaya yardımcı olmuş olmasını umuyorum.
    Bu tür çalışmalar verimli makine öğrenimi araştırmaları için önemli; şu anda bu alan için yapılabilecek en önemli işlerden biri bile olabilir.
    Araştırma, inovasyon hızıyla ilerler; inovasyon da deney çalıştırma süresinin tersiyle hızlanır ve bu açıkça araştırma ya da basit hack amaçlı kodun Kolmogorov karmaşıklığıyla ilişkilidir.
    Bu tür araçların araştırma için ne kadar önemli olduğunu, kişisel olarak bilgi keşfi sürecimi ne kadar hızlandırdığını ne kadar vurgulasam az. Fikirleri birkaç dakika içinde hızla taslak hâline getirip anında sinyal-gürültü oranı yüksek sonuçlar alabilmek, araştırma ilerlemesinin vazgeçilmez bir parçası hâline geldi.
    Bilgi damıtma ve MDL(https://en.wikipedia.org/wiki/Minimum_description_length), mevcut makale gönderim ve inceleme sürecinin teşvik ediyor gibi göründüğü gereksiz süslemeleri, kalabalığı ve aşırı yoğun “geri kalmama” odaklı düşük değerli konu yarışını geri çevirmek için çok önemli bence.
    Son zamanlarda bu sorundan kaçınmak ve biraz daha iyi bir ölçekleme çözümüne gitmek istediğim için, kodu tek dosyalık, kendi kendine yeten kısa gist’ler olan “kod eskizleri” olarak yayımlamaya başladım. Geliştirme süresini kısaltıyor ve kavramı içeren kaba, cilalanmamış ama çalışan kodu insanlara doğrudan ulaştırabiliyor. Şimdiye kadar oldukça iyi işliyor gibi görünüyor ve devam etmek istiyorum.
    Böyle kodları daha fazla görmek istiyorum. Büyük ölçekte veriyle model eğiten araştırmacılar, bilgi yayma biçimlerinde de veri açısından verimli olmalı.

    • 2023, yapay zeka araştırmalarının akıl almaz bir hızla geliştiğini sadece izlemek açısından bile ilginç bir yıldı. ArXiV, PyTorch, GitHub, Hugging Face, özlü açık kaynak Python kodu gibi temel unsurlar bu yeni alanın gelişimini dramatik biçimde hızlandırıyor.
      İnsanlığın belirli bir karmaşıklığa sahip bir şeyi daha önce bu kadar hızlı geliştirmiş olması muhtemelen pek olası değil.
      Benzer hızın görüldüğü yerlerden biri ancak SpaceX; bu yıl da son teknoloji iki roket fırlattı. 2024’te neler çıkacağını merak ediyorum.
    • Küçük bir performans iyileştirme olasılığı var. Burada x_proj için bias olmadığına göre x_proj ve dt_proj ağırlıkları birleştirilebilir gibi görünüyor.
      Ağırlık ayarlama gereksinimleri varsa bu runtime’da basitçe yapılabilir belki; tek bir kernel ve bias’ın sonunda daha hızlı olacağını düşünüyorum. Emin değilim.

  • Orijinal makale tartışması var mıydı merak ediyorum. Kaçırmış gibiyim ama oldukça ilginç.
    “Verimli uygulama eksikliği bellek yetersizliğine ya da gerçekçi olmayan hesaplama gereksinimlerine yol açtığından, SSM olarak da yorumlanabilecek önceki güçlü döngüsel modeller olan RWKV ve RetNet baz çizgilerinin 8k bağlam uzunluğundaki tam sonuçları eksik” kısmını tam anlayamadım.
    RetNet çok bellek kullanmıyor ve parçalar hâlinde ileri geçiş uygulaması kullanılırsa VRAM kullanımı parça boyutuyla sınırlanıyor. Bağlam uzunluğunu test etmenin özü de bu kısım.
    Orijinal Mamba modelini test etmiş biri var mı merak ediyorum. Paralel ileri geçiş modundaki RetNet’e kıyasla eğitim hızı ne düzeyde?

  • Karmaşık bir şeyi yalnızca özü kalacak şekilde sadeleştiren uygulamalar her zaman iyidir.