Rust+WASM के साथ heterogeneous Edge पर तेज़ और portable Llama2 inference चलाना

• Rust+WASM स्टैक AI inference में Python का एक शक्तिशाली विकल्प बन सकता है
  • Python की तुलना में Rust+Wasm ऐप का आकार 1/100, गति 100 गुना तेज़ हो सकती है, और binary code बदले बिना full hardware acceleration के साथ हर जगह सुरक्षित रूप से चलाया जा सकता है
  • Rust, AGI की भाषा है
• अगर Llama2 inference को native speed पर चलाने वाला एक सरल Rust प्रोग्राम लिखें
  • Wasm compile के बाद binary केवल 2MB की रहती है, लेकिन heterogeneous hardware accelerators पर पूरी तरह portable होती है
  • साथ ही Wasm runtime (WasmEdge) cloud environment के लिए एक सुरक्षित execution environment देता है. यह container tools के साथ सहज रूप से काम करता है, जिससे विभिन्न devices पर portable applications को orchestrate और run किया जा सकता है

Python क्यों नहीं?

• llama2 जैसे LLM आम तौर पर Python (जैसे: PyTorch, Tensorflow, JAX) में train किए जाते हैं
• लेकिन AI computing के लगभग 95% हिस्से, यानी inference applications, के लिए Python का उपयोग करना एक बड़ी गलती है
• Python packages में जटिल dependencies होती हैं. इन्हें set up करना और इस्तेमाल करना कठिन है
• Python dependencies बहुत भारी होती हैं. Python या PyTorch के Docker images आम तौर पर कई GB या दर्जनों GB के होते हैं, जो edge server या device पर AI inference चलाते समय खास तौर पर समस्या बनते हैं
• Python बहुत धीमी भाषा है. यह C, C++, Rust जैसी compiled languages की तुलना में अधिकतम 35,000 गुना तक धीमी हो सकती है
• Python धीमी होने के कारण अधिकांश वास्तविक workloads को Python wrapper के नीचे मौजूद native shared libraries पर छोड़ना पड़ता है
  • इसलिए Python inference apps demo के लिए तो बेहतरीन हैं, लेकिन business-specific requirements के अनुसार अंदरूनी बदलाव करना बहुत कठिन है
  • Native libraries पर अधिक निर्भरता और dependency management की जटिलता के कारण, device की विशेष hardware capabilities का उपयोग करते हुए Python AI प्रोग्राम को कई devices पर port करना बहुत कठिन है
• LLM toolchain में आम तौर पर उपयोग होने वाले Python packages अक्सर एक-दूसरे से सीधे टकराते हैं
• LLVM, TensorFlow और Swift भाषा के लिए प्रसिद्ध Chris Lattner ने इस हफ्ते startup podcast में एक शानदार इंटरव्यू दिया, जिसमें उन्होंने बताया कि Python model training के लिए तो बेहतरीन है, लेकिन inference applications के लिए गलत विकल्प क्यों है
  • Expanding AI chip capabilities beyond Nvidia with Modular CEO Chris Lattner

Rust+Wasm के फायदे

• Rust+Wasm स्टैक devices से लेकर edge cloud, on-premise servers और public cloud तक एकीकृत cloud computing infrastructure प्रदान करता है
• यह AI inference applications के लिए Python स्टैक का एक शक्तिशाली विकल्प है, और Elon Musk ने Rust को AGI की भाषा कहा, यह आश्चर्य की बात नहीं है
• अत्यंत हल्का: inference application सभी dependencies सहित केवल 2MB का है. यह सामान्य PyTorch container size के 1% से भी कम है
• बहुत तेज़: preprocessing, tensor operations, postprocessing सहित inference application के हर हिस्से में native C/Rust speed देता है
• portable: heterogeneous hardware acceleration को सपोर्ट करने वाले सभी प्रमुख computing platforms पर वही Wasm bytecode application चल सकता है
• set up, development और deployment आसान: अब जटिल dependencies नहीं हैं. लैपटॉप के standard tools से एक single Wasm file build करके कहीं भी deploy किया जा सकता है
• सुरक्षित और cloud-ready: Wasm runtime को untrusted user code को isolate करने के लिए डिज़ाइन किया गया है. Wasm runtime को container tools से manage किया जा सकता है और cloud-native platforms पर आसानी से deploy किया जा सकता है

Rust inference demo

• 40 lines के Rust code में लिखा गया है और Wasm में compile किया गया है
• Wasm bytecode बनने के बाद इसे WasmEdge runtime को support करने वाले किसी भी device पर deploy किया जा सकता है
• llama.cpp पर आधारित WasmEdge GGML plugin, device के hardware acceleration का अपने-आप उपयोग करके llama2 model चलाता है

अगले कदम

• WasmEdge GGML tooling अभी उपलब्ध है और वास्तव में cloud-native customers द्वारा उपयोग भी किया जा रहा है, लेकिन यह अभी शुरुआती चरण में है
• यदि आप open source project में योगदान देने और भविष्य की LLM inference infrastructure की दिशा तय करने में रुचि रखते हैं, तो संभावित योगदान के कुछ क्षेत्र ये हैं
  • अधिक hardware और OS platforms के लिए GGML plugin जोड़ना: जैसे Linux और Windows पर TPU, ARM NPU और अन्य विशेष AI chips
  • अधिक llama.cpp configurations को support करना: फिलहाल Wasm से GGML plugin तक कुछ configuration options भेजने का समर्थन है, लेकिन लक्ष्य GGML द्वारा दिए गए सभी options को support करना है
  • अन्य Wasm-compatible languages के WASI NN API को support करना: खासकर Go, Zig, Kotlin, JavaScript, C और C++

[gcback]

फ़िलहाल AI में पीछे चल रही Apple के लिए सबसे व्यावहारिक रास्ता शायद यही लगता है कि,

मुख्य लेख में बताए गए जैसे optimized S/W inference engine को iPhone A-chip के NPU से accelerate किया जाए और foundation model फ़ाइल को पूरा का पूरा iPhone में डाल दिया जाए।
हाल में अकादमिक और इंडस्ट्री, दोनों में यह आम राय है कि AI performance में size से ज़्यादा data और fine-tuning महत्वपूर्ण हैं, और Llama2 में fine-tuning performance काफ़ी अच्छी मानी जाती है। इसलिए अच्छी तरह train किए गए मॉडल को लगातार edge device पर deploy करना भी पूरी तरह संभव लगता है। उस नज़रिए से iPhone में सब कुछ ठूंस देने की strategy काफ़ी व्यावहारिक लगती है.

[ragingwind]

क्या WASM सच में इतना लोकप्रिय हो रहा है? लेकिन क्या Python की तुलना Mojo से नहीं करनी चाहिए?

[freedomzero]

सहमत हूँ। Mojo को अभी भी experimental stage में ही माना जाना चाहिए।

[cosine20]

Mojo अभी तक आधिकारिक रूप से रिलीज़ नहीं हुआ है, लेकिन Rust पहले से ही काफ़ी mature language बन चुकी है और real-world काम में कई हिस्सों को replace कर रही है।

[xguru]

GN⁺: heterogeneous edge पर तेज़ और portable Llama2 inference
GN⁺ का सारांश और Hacker News टिप्पणियाँ भी साथ में देखें।