PyTorch की एक फ़ाइल में Mamba का इम्प्लीमेंटेशन

 (github.com/johnma2006)
1 पॉइंट द्वारा GN⁺ 2023-12-21 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • mamba-minimal एक प्रोजेक्ट है जो Mamba को PyTorch की एक ही फ़ाइल में सरल और न्यूनतम रूप में इम्प्लीमेंट करता है
  • इसका लक्ष्य आधिकारिक इम्प्लीमेंटेशन के साथ forward/backward pass में समान numerical output देना है
  • कोड को सरल बनाया गया है और पढ़ने में आसान रखने के लिए comments सहित तैयार किया गया है
  • इसमें आधिकारिक इम्प्लीमेंटेशन की मुख्य optimizations शामिल नहीं हैं, इसलिए यह speed प्रदान नहीं करता, और न ही इसमें उचित parameter initialization शामिल है
  • डेमो state-spaces/mamba-370m और EleutherAI/gpt-neox-20b tokenizer का उपयोग करके prompt completion का उदाहरण चलाता है

प्रोजेक्ट का अवलोकन

  • mamba-minimal Mamba का PyTorch की एक फ़ाइल में बना एक सरल न्यूनतम इम्प्लीमेंटेशन है
  • लक्ष्य है कि आधिकारिक इम्प्लीमेंटेशन जैसा व्यवहार अधिक पढ़ने योग्य कोड में दिखाया जाए
  • मुख्य विशेषताएँ:
    • आधिकारिक इम्प्लीमेंटेशन के forward pass और backward pass में समकक्ष numerical output
    • सरल किया गया कोड
    • पढ़ने में आसान और comments वाला इम्प्लीमेंटेशन

इसमें शामिल नहीं है

  • speed इसका लक्ष्य नहीं है
    • आधिकारिक इम्प्लीमेंटेशन काफ़ी अधिक optimized है
    • यह optimization Mamba पेपर के मुख्य योगदानों में शामिल है
    • यह इम्प्लीमेंटेशन readability के लिए अधिकांश भागों को सरल रखता है
  • उचित parameter initialization शामिल नहीं है
    • इसे ऐसी चीज़ के रूप में बताया गया है जिसे readability से समझौता किए बिना जोड़ा जा सकता है

डेमो उपयोग उदाहरण

  • demo.ipynb में prompt completion का उदाहरण देखा जा सकता है
  • उदाहरण model.Mamba और Hugging Face transformers के AutoTokenizer का उपयोग करता है
  • उपयोग किए गए model और tokenizer:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • उदाहरण prompt Mamba is the है, और generated output में Mamba को एक venomous snake के रूप में वर्णित करने वाला वाक्य शामिल है

संदर्भ सामग्री

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2023-12-21
Hacker News की राय

  • पहले मैंने एक सहकर्मी के साथ मिलकर एक library बनाई थी, जिसमें shared model code का ज़्यादातर हिस्सा अलग कर दिया गया था; इसका इस्तेमाल करने पर कई models को Python import और comments को छोड़कर लगभग 100 lines में implement किया जा सकता है
    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, PyTorch flash attention जैसी सुविधाएँ भी support करती है

    • मैं इस library को ज़रूर देखूँगा। जानना चाहूँगा कि आपने xformers भी देखा है या नहीं
      xformers भी मिलती-जुलती समस्या को address करता है, लेकिन Triton का इस्तेमाल करके high-performance Transformer modules देने पर उसका focus ज़्यादा है। हालांकि library के सिर्फ कुछ specific components निकालकर इस्तेमाल करना आसान नहीं था, और runtime errors लगातार आ रहे थे, इसलिए फिलहाल उसे टाल दिया। मैं BERT architecture पर आधारित कुछ बना रहा हूँ, इसलिए इसे reference के तौर पर देखूँगा
    • यह library देखकर मैं प्रभावित हूँ। Hugging Face की implementation मुझे ख़ास पसंद नहीं थी, लेकिन यह abstraction level के लिहाज़ से बिलकुल सही, सुंदर API जैसा दिखता है
      अगले project में इसे आज़माने का सोच रहा हूँ

  • मूल Mamba code में speed optimization और अन्य चीज़ें बहुत हैं, इसलिए उसे सीधे समझना मुश्किल है; यह implementation सीखने में मददगार लगती है
    जब token-by-token inference करते हैं, तो सब कुछ काफ़ी सरल हो जाता है। मेरी अपनी बनाई Mamba inference implementation भी है: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran? आपने Fortran क्यों इस्तेमाल किया, यह जानने की उत्सुकता है
      मुझे पता है कि यह लंबे समय से verified scientific computing code का आधार रहा है और अक्सर PyTorch या Numpy जैसी libraries से wrap करके इस्तेमाल किया जाता है, लेकिन आजकल यह popular language नहीं है। इसे चुनने की वजह जानना चाहूँगा

  • Mamba के बारे में कुछ बातें ऐसी हैं जिन्हें मैं चाहता/चाहती हूँ कि कोई non-ML researcher भी समझ सके इस तरह समझाए

    1. Transformer से आगे state-space models की कुल मिलाकर insight क्या है
    2. Mamba ने S4, H3, Monarch जैसे पिछले कामों की तुलना में कौन-सी incremental innovations की हैं, जिनसे वह ज़्यादा सफल या दिलचस्प बना
    3. context length की sub-quadratic scalability के अलावा इसका क्या मतलब है। जैसे अगर मुझे 100k tokens से ज़्यादा की context length में दिलचस्पी नहीं है, तो क्या समान आकार के model और dataset पर Mamba की training compute efficiency बेहतर होने की संभावना है?
    • paper authors की तुलना में मेरी बुद्धि काफ़ी कम है, लेकिन फिर भी मैंने समझने की कोशिश की। मैंने computer science पढ़ी है और basic control theory व undergraduate स्तर की discrete-time systems intuition है, लेकिन लगता है इस paper को ठीक से समझने के लिए state-space models कहीं ज़्यादा पढ़ने होंगे
      Mamba की core insight state-space models की एक पुरानी समस्या हल करने में है। state-space models input context को compress करने में अच्छे होते हैं, लेकिन input को hidden state में compress करने की प्रक्रिया में वह जानकारी मिट जाती है जो Transformer की तरह context को प्रभावी ढंग से इस्तेमाल करने के लिए ज़रूरी होती है
      समाधान यह है कि paper जिसे selection mechanism कहता है, उसे बनाया जाए। यह mechanism input-dependent है, इसलिए input बदलने पर model हर step के output को adjust कर सकता है। इसके लिए कुछ state-space variables को input-invariant के बजाय input-dependent बनाया जाता है, और हर time step के input को state-space variables में project करने के लिए linear layers आदि जोड़ी जाती हैं
      लेकिन state-space variables को input-dependent बनाने से computational overhead आता है। इसे आधुनिक GPU memory structure का अधिकतम उपयोग करने वाले hardware-aware algorithm से हल किया गया, ताकि HBM के अंदर-बाहर data move करने से जितना हो सके बचा जा सके
      Tri Dao वही व्यक्ति हैं जिन्होंने Flash Attention बनाया था, और वह भी Transformer में hardware को ज़्यादा efficiently इस्तेमाल करने का तरीका था। यह सचमुच उनकी विशेषज्ञता का क्षेत्र है
    • Attention context length के साथ quadratic रूप से बढ़ता है, gating वाले recurrent neural networks (LSTM, GRU आदि) linear होते हैं, और ये नए architectures भी linear हैं। शुरुआती recurrent neural networks exploding gradients से बचने के लिए gating इस्तेमाल करते थे, लेकिन नए approaches stability guarantee करने वाली dynamical systems theory का उपयोग करते हैं, ताकि gating दोनों समस्याएँ एक साथ हल करने के बजाय memory पर focus कर सके
      NeurIPS 2023 से ठीक पहले आए Mamba और Based में multi-query associative recall (MQAR) और multi-head Attention से प्रेरित gating/selection की data-dependency शामिल थी। ये दोनों वे key elements थे जो Hyena और पहले के state-space architectures में नहीं थे, और नए models associative recall tasks में Attention जितने अच्छे हो गए हैं; lookup के अलावा दूसरे tasks में शायद Attention से थोड़ा बेहतर होने की संभावना भी दिखती है
      बेशक Mamba की बड़ी detail उसका efficient CUDA implementation है। उसके बिना, जहाँ Transformer पहले से fit बैठता है उन tasks में इस architecture का महत्व कम हो सकता है
      context length की बहुत चिंता न भी करें, तब भी कई नए areas खुलते हैं। DNA sequence analysis लंबी dependencies वाला linear task है, और images, videos, high-dimensional information को token streams के रूप में देखने का तरीका भी सोचा जा सकता है। यह पुराने CRT monitors की तरह pixels को scan करने जैसा है
      AI के शुरुआती सपनों में से एक यह था कि environment के साथ लगातार interact करने वाले agent की single learning trajectory लगातार evolve होती रहे, और ऐसे infinite context length models शायद उस सपने को आसान बना सकते हैं
      हालांकि अभी के लिए, ऐसे models के महत्वपूर्ण real-world tasks पर downstream applications, Attention-based mature applications की तुलना में आम तौर पर कम validated और tuned हैं। पुराने recurrent neural networks से analogy कुछ हद तक मदद करती है, लेकिन पिछले 5 वर्षों में लोग Attention और Transformer के लिए बहुत ज़्यादा specialized हो गए हैं, इसलिए Transformer की inertia बहुत बड़ी है
    • समान आकार के model और dataset पर Mamba ज़्यादा compute-efficient तरीके से train हो सकता है या नहीं, यह मैं भी जानना चाहूँगा/चाहूँगी
      original paper बताता है कि parameters transform होने के बाद model को linear recurrence या global convolution, दोनों तरीकों से compute किया जा सकता है। सामान्य तौर पर training में, जहाँ पूरी input sequence पहले से देखी जा सकती है, parallelize करना आसान होने के कारण convolution mode इस्तेमाल होता है; और autoregressive inference में, जहाँ input एक time step करके देखा जाता है, efficient recurrent mode में switch किया जाता है
      इसलिए training RetNet के parallel forward-pass mode की तरह parallelizable है। basic inference लंबे से लंबा context पाने के लिए recurrent mode में किया जाता है, और chunking न होने के कारण inference के दौरान RAM और VRAM कितना खाएगा, यह आंकना मुश्किल है
    • यह video शायद ठीक वही है जिसे आप ढूँढ रहे थे
      यह paper समझाते हुए भी big picture में वह कहाँ fit होता है, इस पर काफ़ी context देता है। इसकी progression सुनना काफ़ी दिलचस्प है
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • मेरी जानकारी में Mamba मूल रूप से state-space model research की उसी धारा का विस्तार है जिसे long convolution कहा जा सकता है
      quadratic Attention के बजाय, जहाँ calculate किया जाता है कि हर token बाकी सभी tokens पर कितना ध्यान देता है, किसी तरह input जितनी लंबाई का long convolution kernel calculate किया जाता है और फिर conv1d apply किया जाता है
      मेरी सीमित समझ के हिसाब से इसका FFT apply करने, matrix multiplication करने, और फिर IFFT से वापस लौटने से थोड़ा संबंध है। पता है कि यह काम करता है, लेकिन धीमा है। FFT compute करने के कई तरीके हैं, और उनमें से एक butterfly matrix है। शायद यह सिर्फ़ approximation है, लेकिन काफी अच्छा है और current hardware पर बहुत fast और efficient लगता है
      quadratic complexity सुनने में खराब लगती है, लेकिन असल में hardware constraints के कारण sub-quadratic algorithms कई बार ज़्यादा slow होते हैं। इसलिए state-space models से बड़ी उम्मीदें होने के बावजूद यह कहना आसान नहीं है कि Llama खत्म हो गया। Mamba scale बढ़ाने पर भी अच्छा चलेगा या नहीं, यह भी अभी नहीं पता, और यह जानने के लिए सचमुच training पर millions of dollars खर्च करने होंगे। फिर भी मैं optimistic हूँ
      sub-quadratic family का एक और दिलचस्प model RWKV है। देखने लायक है, लेकिन शायद podcast में इसे पहले ही cover किया गया होगा
      मैंने self-study की है और paper भी पहले बस roughly skim किया था, इसलिए बहुत कुछ गलत हो सकता है। साथ ही Attention में आम तौर पर KV cache होता है जो performance में बहुत मदद करता है, लेकिन मुझे लगता है Mamba में वह नहीं किया जा सकता

  • “Mamba दुनिया का सबसे लंबा विषैला सांप है, जिसकी अनुमानित लंबाई 150m से ज़्यादा है” वाले वाक्य पर हंसी आ गई
    फिर भी यह वाकई शानदार था, और arXiv पेपर का संदर्भ देने से मेरे जैसे लोग, जो पेपर को सीधे समझने के बजाय ऐसे लेख पढ़ते हैं, अंदर की चीज़ों की थोड़ी झलक पा सके—यह अच्छा लगा

    • Mamba नाम अच्छा है। [S]elective [S]tructured [S]tate [S]pace [S]equence models होने से sSSSS बनता है, जो सांप की आवाज़ जैसा लगता है
    • मुझे लगा था सबसे लंबा विषैला सांप किंग कोबरा होता है। एक साधारण Google खोज में भी यही निकला
      अगर बाद में उस वाक्य पर सुधार जारी करना पड़े, तो मज़ेदार होगा

  • मैंने सोचा था कि एल्गोरिदम का मूल parallel prefix scan होगा। मुझे लगता है 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)

  • यह मूर्खतापूर्ण सवाल हो सकता है, लेकिन Hugging Face पर मौजूद Mamba model को train करना कितना कठिन है, यह जानना चाहता हूं
    सबसे बड़ा model 2.8B का दिखता है; The Pile जैसे dataset पर train करने के लिए कितने GPU चाहिए होंगे और कितना समय लगेगा?

    • यह एक शानदार सवाल है, जिसे मैं भी जानना चाहता हूं। जवाब यह लगता है कि समान आकार के Transformer की तुलना में यह काफ़ी तेज़ है, और अंतिम परिणाम भी लगभग सभी benchmarks पर Transformer से बेहतर score कर सकते हैं
      inference भी RAM की आधी खपत के साथ 3~5 गुना तेज़ होने की उम्मीद है

  • आधिकारिक CUDA version को समझने की कोशिश की थी, लेकिन पहली कोशिश नाकाम होने के बाद आखिरकार छोड़ दिया; यह implementation कहीं बेहतर लगती है

  • एक और single-file PyTorch implementation, सच में शानदार। उम्मीद है कि पहले के hlb-CIFAR10 और संबंधित projects, और minGPT या DawnBench जैसे पूर्व प्रभावों ने simple single-file format को थोड़ा भी आगे बढ़ाने में मदद की होगी
    ऐसा काम efficient machine learning research के लिए महत्वपूर्ण है, और इस क्षेत्र के लिए अभी किए जा सकने वाले सबसे महत्वपूर्ण कामों में से एक हो सकता है
    research innovation की गति से आगे बढ़ती है, और innovation प्रयोगों के execution time के inverse के अनुपात में तेज़ होती है; यह research या simple hacking उद्देश्यों के लिए code की Kolmogorov complexity से स्पष्ट रूप से जुड़ा है
    ऐसे tools research के लिए कितने महत्वपूर्ण हैं, और व्यक्तिगत रूप से उन्होंने knowledge discovery process को कितना तेज़ किया है, इसे जितना भी कहें कम है। किसी idea को कुछ मिनटों में जल्दी sketch करना और तुरंत high signal-to-noise ratio वाले results पाना research progress के लिए अनिवार्य तत्व बन गया है
    knowledge distillation और MDL(https://en.wikipedia.org/wiki/Minimum_description_length) मेरे हिसाब से अनावश्यक सजावट, clutter, और अत्यधिक घने “पीछे न छूटने” वाले कम-मूल्य topic competition को उलटने के लिए बहुत महत्वपूर्ण हैं, जिन्हें मौजूदा paper submission-review process बढ़ावा देता दिखता है
    हाल में, इस समस्या से बचने और थोड़ा बेहतर scaling solution की ओर जाने के लिए, मैंने code को 1-file self-contained छोटी gist वाले “code sketch” के रूप में release करना शुरू किया है। इससे development time घटता है, और concepts को समेटे rough, unpolished working code सीधे लोगों तक पहुंच सकता है। अब तक यह काफ़ी अच्छा काम करता दिख रहा है और मैं इसे जारी रखना चाहता हूं
    ऐसा code और ज़्यादा देखना चाहता हूं। अगर बड़े scale पर data सीखने वाले researchers हैं, तो उन्हें information फैलाने के तरीके में भी data-efficient होना चाहिए

    • 2023 सिर्फ़ यह देखने के लिए भी रोमांचक साल था कि AI research कितनी अविश्वसनीय रफ़्तार से आगे बढ़ रही है। ArXiV, PyTorch, GitHub, Hugging Face, concise open-source Python code जैसे आधारभूत तत्व इस नए क्षेत्र की प्रगति को नाटकीय रूप से तेज़ कर रहे हैं
      शायद मानवता ने इतनी complexity वाली किसी चीज़ को इतनी तेज़ी से पहले कभी विकसित नहीं किया
      मिलती-जुलती रफ़्तार जहां दिखती है, वह शायद SpaceX है, जिसने इस साल भी दो cutting-edge rockets launch किए। 2024 में क्या निकलेगा, यह देखना दिलचस्प होगा
    • performance में छोटी improvement की संभावना है। यहां x_proj में bias नहीं है, इसलिए लगता है कि x_proj और dt_proj weights को combine किया जा सकता है
      अगर weight adjustment की requirement हो तो runtime पर शायद इसे सरलता से किया जा सकता है, और single kernel और bias अंततः तेज़ होंगे। पक्का नहीं

  • सोच रहा हूं कि original paper पर discussion हुआ था या नहीं। लगता है मुझसे छूट गया, लेकिन यह काफ़ी दिलचस्प है
    “efficient implementation की कमी के कारण memory shortage या अव्यावहारिक compute requirement पैदा हुई, इसलिए SSM के रूप में भी व्याख्यायित किए जा सकने वाले पहले के मजबूत recurrent models RWKV और RetNet baselines के context length 8k के पूरे results गायब हैं” वाला हिस्सा मुझे ठीक से समझ नहीं आया
    RetNet ज़्यादा memory नहीं खाता, और chunk-wise forward implementation इस्तेमाल करने पर VRAM usage chunk size तक सीमित रहता है। context length test करने में यही बात मुख्य है
    सोच रहा हूं कि क्या किसी ने original Mamba model test किया है। parallel forward mode वाले RetNet की तुलना में training speed कैसी होगी?

  • जटिल चीज़ को सिर्फ़ core बचाकर साफ़ करने वाला implementation हमेशा अच्छा लगता है