Google च्या Gemma-4 31B मॉडेलला AWS Inferentia2 inf2.24xlarge वर पोर्ट केल्यावर CPU रेफरन्सशी तंतोतंत (token-for-token) मॅच मिळाले—तरीही प्रत्येक तयार झालेले वाक्य अर्थहीन (gibberish) होते. "मॅचिंग" आणि "काम करणे" यातील ही दरी आता Amazon च्या कस्टम इन्फरन्स चिप्सवर मोठे LLMs वापरण्याचा प्रयत्न करणाऱ्यांसाठी एक इशारा आहे.

टोकन-बाय-टोकन मॅच पुरेसा का नाही

डेव्हलपरने Inferentia डिव्हाइसमधून मिळणारे प्रत्येक आउटपुट टोकन आणि मॉडेलच्या CPU रनमधून तयार झालेले टोकन यांची तुलना केली. दोन्ही स्ट्रीम्स एकसारख्या होत्या, त्यामुळे हार्डवेअरने रेफरन्स इम्प्लिमेंटेशनची हुबेहूब नक्कल केली आहे असे वाटले. प्रत्यक्षात, दोन्ही स्ट्रीम्सनी मॉडेलमध्ये चुकीचा प्रॉम्प्ट (malformed prompt) पाठवला होता, कारण मॉडेलमधील चॅट टेम्पलेट काढले गेले होते आणि चुकीचे टर्न मार्कर्स (turn markers) दिले गेले होते. टेम्पलेट नसल्यामुळे मॉडेल अनंत लूपमध्ये (infinite loop) अडकले आणि अर्थहीन मजकूर बाहेर काढू लागले. हार्डवेअरने आपले काम केले—त्यांनी रेफरन्स कोडमधील बगची हुबेहूब नक्कल केली.

धडा साधा आहे: SEQ_MATCH (sequential token equality) म्हणजे अचूकता (correctness) नव्हे. जर रेफरन्स इम्प्लिमेंटेशनमध्येच दोष असेल, तर हार्डवेअरची हुबेहूब प्रतिकृती देखील तोच दोष वारसा म्हणून घेते. व्हॅलिडेशन केवळ टोकन-लेव्हल पॅरिटीपुरते मर्यादित नसावे; त्यासाठी योग्यरित्या फॉरमॅट केलेल्या इनपुट्ससह एंड-टू-एंड फंक्शनल चेकची आवश्यकता आहे.

पॅरामीटर्सच्या रूपात लपलेले बफर्स (Buffers)

लोडिंग फेज दरम्यान, मॉडेल लोडरने layer_scalar नावाचा घटक वगळला. PyTorch मॉडेल डेफिनेशनमध्ये कोडने या ऑब्जेक्टला parameter ऐवजी buffer म्हणून नोंदवले होते. बफर्स हे स्टॅटिक टेन्सर (static tensors) असतात जे ट्रेनिंग दरम्यान अपडेट होत नाहीत आणि अनेक लोडर्स Neuron-सुसंगत फॉरमॅटमध्ये रूपांतरित करताना त्यांना दुर्लक्षित करतात. हे वगळल्यामुळे अनेक लेयर्सचे स्केलिंग फॅक्टर्स त्यांच्या डिफॉल्ट व्हॅल्यूवर राहिले, ज्यामुळे संपूर्ण नेटवर्कमधील गणिती प्रक्रिया विस्कळीत झाली. कोणताही एरर आला नाही; मॉडेल कंपाईल झाले आणि इन्फरन्स पाइपलाइन चालली, परंतु गणिती निकाल चुकीचे होते.

Inferentia वर मोठे मॉडेल्स नेणाऱ्या प्रत्येकाने प्रत्येक नॉन-पॅरामीटर टेन्सरची तपासणी (audit) करावी. जरी एखादा टेन्सर शिकण्यासाठी (to be learned) नसेल, तरीही तो योग्य फॉरवर्ड-पास कम्प्युटेशनसाठी आवश्यक असू शकतो. बफर समाविष्ट असल्याची मॅन्युअली पडताळणी केल्यास 'सायलेंट स्केल एरर्स' टाळता येतील, ज्यांचे निदान करणे कठीण असते.

स्पॉट-इन्स्टन्सची अस्थिरता आणि ३९ मिनिटांचा कंपाईल वेळ

स्पॉट इन्स्टन्सवर ३१-बिलियन पॅरामीटरचे मॉडेल चालवणे स्वस्त वाटते, परंतु या बचतीसोबत अनपेक्षित रिकलेम इव्हेंट्स (reclaim events) येतात. डेव्हलपरचा कंपाईल वेळ—मॉडेलला Neuron-सुसंगत कोडमध्ये रूपांतरित करण्यासाठी लागणारे सुमारे ३९ मिनिटे—AWS ने इन्स्टन्स परत घेतल्यावर वाया गेले. व्यत्यय टाळण्यासाठी त्यांनी तीन स्तरांची सुरक्षा यंत्रणा तयार केली:

  • ModelBuilder ने मेमरी वापर ३८४ GB होस्ट मर्यादेत ठेवला, ज्यामुळे क्रॅश होऊन पुन्हा स्टार्टअप करण्याची वेळ आली नाही.
  • रॉ वेट फाइल्स (raw weight files) आणि कंपाईल केलेल्या “neffs” (Neuron executable files) या दोन्हीची त्वरित S3 मिररिंग केल्यामुळे, नवीन इन्स्टन्स नेमक्या त्याच ठिकाणाहून काम सुरू करू शकला जिथे आधीचा इन्स्टन्स थांबला होता.
  • एका multi-region poller ने उपलब्ध स्पॉट कॅपॅसिटीसाठी AWS रीजन्स स्कॅन केले आणि क्षमता उपलब्ध होताच नवीन इन्स्टन्स सुरू केला.

या पावलांमुळे एक नाजूक, सिंगल-पॉइंट कंपाईल प्रक्रिया एका लवचिक (resilient) पाइपलाइनमध्ये बदलली, जी स्पॉट मार्केटमधील चढ-उतारांमध्येही टिकून राहते.

मिश्र अटेंशन लेआउट्ससह शार्डिंगमधील अडचणी

Gemma-4 31B मध्ये दोन अटेंशन कॉन्फिगरेशन वापरले जातात. काही लेयर्समध्ये चार की-व्हॅल्यू (KV) हेड्स वापरले जातात, तर इतरांमध्ये वेगळी संख्या असते. जेव्हा एखाद्या लेयरमधील KV हेडची संख्या आठ समांतर रँक्समध्ये (parallel ranks) व्यवस्थित विभागली जात नाही, तेव्हा मॉडेल आठ रँक्समध्ये समान रीतीने विभाजित करणे (sharding) अपयशी ठरते. ४-हेड असलेल्या लेयरला आठ रँक्समध्ये विभाजित करण्याचा प्रयत्न केल्यास प्रत्येक रँकला अर्धा हेड हाताळावा लागेल—ही एक गणितीदृष्ट्या अशक्य गोष्ट आहे, ज्यामुळे शेप मिसमॅच (shape mismatches) आणि रनटाइम एरर्स येतात.

यावर उपाय म्हणजे ग्लोबली-शार्ड केलेल्या लेयर्सची (ज्यांची हेड संख्या सुसंगत आहे) सर्व रँक्समध्ये पुनरावृत्ती (replicate) करणे आणि फक्त अशा “स्लाइडिंग” लेयर्सचे शार्डिंग करणे ज्यांच्या हेड संख्या समान विभागणीला परवानगी देतात. या हायब्रिड स्ट्रॅटेजीमुळे KV हेड्सचे अवैध विभाजन टाळले गेले आणि टेन्सर-पॅरलल कार्यक्षमता कायम ठेवली गेली, ज्यामुळे पूर्वीच्या प्रयत्नांमध्ये येणाऱ्या टेन्सर-पॅरललायझेशन एरर्स दूर झाल्या.

निष्कर्ष (Takeaway)

एका विशाल LLM ला Inferentia वर पोर्ट करणे म्हणजे केवळ 'कंपाईल आणि रन' करण्याचा व्यायाम नाही. त्यासाठी टोकन इक्वालिटीच्या पलीकडे जाऊन कठोर फंक्शनल टेस्टिंग, प्रत्येक टेन्सर—पॅरामीटर किंवा बफर—योग्यरित्या हाताळला गेला आहे याची सूक्ष्म पडताळणी आणि स्पॉट-इन्स्टन्स रिकलेमेशनचा विचार करणारी डिप्लॉयमेंट स्ट्रॅटेजी आवश्यक आहे. शेवटी, शार्डिंगने मॉडेलच्या अंतर्गत अटेंशन भूमितीचा (attention geometry) आदर केला पाहिजे; अन्यथा, वेग देण्याचे आश्वासन देणारे पॅरललिझम हे शांत अपयशाचे (silent failure) कारण ठरू शकते.