Apple MLX में CUDA support जोड़ा जा रहा है

• MLX Apple Silicon आधारित मशीन लर्निंग शोधकर्ताओं के लिए एक array framework है, जो NumPy और PyTorch जैसी usability प्रदान करता है
• MLX में CUDA backend जोड़ने का काम चल रहा है
• मुख्य उद्देश्य हैं unified memory support और NVIDIA hardware का व्यापक उपयोग
• अभी केवल tutorial examples ही चल पा रहे हैं, और build व test को Ubuntu 22.04 + CUDA 11.6 environment में verify किया गया है
• शुरुआती performance issues और bottleneck सुधार की प्रक्रिया से गुजरते हुए optimization और refactoring बार-बार किए जा रहे हैं
• MLX का CUDA backend Apple के समर्थन से विकसित किया जा रहा है, और unified memory व NVIDIA hardware support के जरिए Mac पर development और बड़े computing environment में deployment के दौरान एक consistent experience देना इसका लक्ष्य है

---

MLX framework का परिचय

• MLX Apple की machine learning research team द्वारा विकसित मशीन लर्निंग के लिए array framework है
• इसे engineers और researchers को प्रभावी ढंग से experiment करने और ideas को तेज़ी से validate करने में मदद देने के लिए बनाया गया है
• इसकी खासियत Apple Silicon support रही है, लेकिन अब CUDA backend जोड़कर NVIDIA GPU environment में भी इसकी उपयोग-क्षमता बढ़ाने की कोशिश की जा रही है

मुख्य विशेषताएँ

• परिचित API: Python API NumPy जैसा है, जबकि upper-level packages mlx.nn, mlx.optimizers PyTorch जैसी interface प्रदान करते हैं
  • C++, C, Swift APIs भी उपलब्ध हैं और Python API के साथ consistent usability सुनिश्चित करते हैं
• Composable function transforms: इसमें automatic differentiation, automatic vectorization, computing graph optimization जैसी सुविधाएँ शामिल हैं
• Lazy evaluation: कई operations को एक साथ बाँधकर केवल आवश्यकता होने पर ही result compute किया जाता है, जिससे resources का अधिक कुशल उपयोग संभव होता है
• Dynamic graph generation: input data का shape बदलने पर भी धीमी compilation के बिना तुरंत execution संभव है, जिससे development और debugging आसान होती है
• Multi-device और unified memory model:
  • पहले यह केवल CPU और Apple GPU को support करता था, लेकिन CUDA backend के साथ NVIDIA GPU environment support की दिशा में काम चल रहा है
  • Unified memory structure अपनाने से memory copy के बिना सभी supported devices पर एक ही object पर operation संभव होता है
• सरल और extensible design: researchers framework को आसानी से extend और improve कर सकते हैं

अन्य frameworks से अंतर और प्रेरणा

• इसकी design inspiration NumPy, PyTorch, Jax, ArrayFire आदि से ली गई है
• विशेष रूप से यह unified memory model, सरल interface, और immediate execution/debugging environment पर ज़ोर देता है

प्रमुख उदाहरण और उपयोग क्षेत्र

• MLX Examples रेपो में विभिन्न practice code उपलब्ध हैं
  • Transformer language model training
  • LLaMA आधारित large-scale text generation और LoRA का उपयोग करके fine-tuning
  • Stable Diffusion से image generation
  • OpenAI Whisper आधारित speech recognition सहित modern model implementations के उदाहरण

MLX का महत्व और CUDA support

• Apple Silicon-केंद्रित framework रहा MLX अब CUDA backend के जुड़ने से NVIDIA GPU environment में भी अधिक सामान्य-purpose उपयोग की संभावना बढ़ा रहा है
• नवीनतम CUDA support के साथ, यह Apple और NVIDIA दोनों hardware पर latest research और industrial use के लिए उपयुक्त framework के रूप में स्थापित हो सकता है

[gosarinamu]

अगर macOS में सिर्फ NVIDIA GPU इस्तेमाल करने दिया जाए, तो बस वही काफ़ी होगा... हाह.

[yangeok]

अच्छी खबर है haha उम्मीद है जल्दी cuda सपोर्ट आ जाए ताकि Mac पर भी हाई-स्पीड ट्रेनिंग संभव हो सके~!

[GN⁺]

Hacker News राय

• मैं बस यह पक्का करना चाहता हूँ कि मेरी समझ सही है: 1, MLX से बना प्रोग्राम CUDA-सपोर्टेड चिप्स पर चल सकता है। लेकिन 2, CUDA प्रोग्राम का Apple Silicon पर चलना संभव नहीं है। क्या इसकी वजह यह है कि 2 कॉपीराइट उल्लंघन होगा, खासकर NVidia की मशहूर एंट्री बैरियर की वजह से?
  • 1 सही है। इससे डेवलपर अपेक्षाकृत कम-शक्ति वाले Apple डिवाइसों (UMA सहित) पर डेवलपमेंट कर सकते हैं, और तैयार कोड को Nvidia के अपेक्षाकृत अधिक-शक्ति वाले सिस्टम पर डिप्लॉय कर सकते हैं। कई कारणों से यह उपयोगी है
  • 2 कॉपीराइट उल्लंघन नहीं है। API का पुनः-इम्प्लीमेंटेशन किया जा सकता है
  • 2 के बारे में, मुझे नहीं लगता कि यह सच में सही बात है। AMD का HIP भी 17~18 साल के आसपास OpenCL छोड़कर यही काम कर रहा है
  • मुझे 3 चाहिए। अगर Apple Silicon से NVIDIA GPU कनेक्ट करके CUDA इस्तेमाल किया जा सके तो अच्छा होगा। Apple Silicon, unified memory, GPU, और CUDA को PyTorch, JAX, या TensorFlow के साथ एक साथ इस्तेमाल करने की कल्पना है, लेकिन मैंने अभी तक MLX को ठीक से आज़माया नहीं है
  • 2 के न होने की बड़ी वजह यह है कि वह कहीं ज़्यादा कठिन है
• अगर आपने सोचा, "रुको, Apple प्लेटफ़ॉर्म पर तो official CUDA support नहीं है?", तो यह भी देखना चाहिए कि यह patch set "CUDA 12 और SM 7.0(Volta) या उससे ऊपर वाले Linux प्लेटफ़ॉर्म" को भी सपोर्ट करता है
  इंस्टॉलेशन गाइड
• इस प्रोजेक्ट को Electron आदि के निर्माता के रूप में मशहूर zcbenz लीड कर रहे हैं
  zcbenz परिचय
• MLX की मुख्य विशेषताओं में से एक unified memory architecture (UMA) का उपयोग है, और मुझे जिज्ञासा है कि यह कैसे काम करता है
  रिपॉज़िटरी readme के bullet देखें: MLX GitHub
  अगर इसे सभी UMA APU तक, निर्माता की परवाह किए बिना, बढ़ाया जाए तो मज़ेदार होगा, लेकिन discrete GPU पर तरीका अलग होना चाहिए या नहीं, यह सवाल है
  PR कमेंट पढ़कर लगता है कि CUDA भी सीधे UMA API को सपोर्ट करता है और ज़रूरत पड़ने पर कॉपी को पारदर्शी तरीके से संभालता है
  • मेरे अनुभव में prefetch फीचर की कमी के कारण copy wait की वजह से memory bottleneck बहुत गंभीर हो जाता है। जब पूरा dataset VRAM में फिट हो जाता है तो manual prefetch ठीक काम करता है, लेकिन मेरे application (ML training) में performance drop इतना बड़ा था कि मुझे सीधे streaming load पर स्विच करना पड़ा
• थोड़ा अलग विषय है, लेकिन MLX पर काम करने वाले कई लोग शायद Apple से आधिकारिक रूप से जुड़े नहीं लगते। उदाहरण के लिए prince_canuma का Twitter देखें
  इससे जिज्ञासा होती है कि क्या Apple पीछे से समर्थन दे रहा है लेकिन आधिकारिक रूप से सामने नहीं आ रहा। टीम के ट्रांसफ़र होने की अफवाहें भी थीं
• PR विवरण में लिखा है, "CUDA backend जोड़ने का ongoing effort"
  यानी इसका मतलब है कि MLX कोड x86 + GeForce हार्डवेयर पर build और run किया जा सकता है, उल्टा नहीं
• "अगर Mac पर कोड लिखकर/टेस्ट करके supercomputer पर deploy किया जा सके तो यह अच्छा developer experience होगा" जैसी बात से यह सवाल उठता है कि क्या अब MLX को Linux पर इस्तेमाल किया जा सकता है
  मैंने खुद टेस्ट किया, और सिर्फ Python 3.12 वर्ज़न ही pypi पर अपलोड है
  MLX-CUDA PYPI पेज
• क्या यह Strix Halo के प्रभाव की वजह से है? मेरे वर्क कंप्यूटर का बजट बच गया था, तो मैंने EVO-X2 खरीदा, और यह midrange znver5 EPYC मशीनों के बराबर या उनसे बेहतर निकला। यह ज़्यादातर EC2 या GCE इंस्टेंसों से, यहाँ तक कि NVMe के मामले में भी, काफ़ी आगे है। इस हार्डवेयर की कीमत 1800 है
  हाल में YouTuber रिव्यू भी आ रहे हैं कि DGX Spark, Strix Halo की तुलना में price-to-performance के हिसाब से कमज़ोर है। कमी यह है कि ROCm अभी भी कमज़ोर है (हालाँकि जल्द बेहतर होने की उम्मीद है)। अगर Apple डिवाइसों पर CUDA पूरी तरह चलने लगे, तो Strix चाहे जितना सस्ता और अच्छा हो, फिर भी यह गंभीरता से विचार करने लायक विकल्प होगा
  • Strix Halo दरअसल AMD Ryzen AI Max+ 395 ही है। यह Framework Desktop और कई mini PC में आ रहा है। इस चिप की memory bandwidth 200GB/s है, जो सामान्य x86 प्लेटफ़ॉर्म के लिए बहुत अच्छी है, लेकिन Nvidia GPU (जैसे 5090 का 1792GB/s) या Apple M3 Ultra (800GB/s) की तुलना में बहुत कम है। price/performance शानदार है, लेकिन LLM जैसे memory-intensive कामों के लिए मुझे लगता है कि यह अभी बस न्यूनतम स्वीकार्य प्रदर्शन की सीमा में प्रवेश कर पाया है
  • PR विवरण देखने से साफ़ पता चलता है कि target cloud cluster training है
  • AMD GPU की neural net performance के Nvidia की बराबरी करने की बात पिछले 10 साल से सुन रहे हैं, लेकिन यह अभी तक सच नहीं हुई है
  • m4 Mac mini की तुलना में यह कैसा है, यह जानने की जिज्ञासा है
• मैं MLX मॉडल से बहुत प्रभावित हूँ। अगर घर के सभी लोगों के लिए local model खोल भी दूँ, तो Nvidia कंप्यूटर की तरह आग लगने के जोखिम की चिंता नहीं करनी पड़ेगी। उम्मीद है Apple Silicon, Nvidia चिप कैंप का एक गंभीर प्रतिस्पर्धी बने। यह भी जिज्ञासा है कि कहीं CUDA support embrace, extend, extinguish (EEE) जैसी रणनीति तो नहीं है
• Apple भविष्य में m-सीरीज़ चिप-आधारित datacenter बनाकर app development, testing, और बाहरी सेवाओं की hosting आदि में उसका उपयोग करने की योजना रखता है