Google के Gemma-4 31B मॉडल को AWS Inferentia2 inf2.24xlarge पर पोर्ट करने से CPU रेफरेंस के साथ एकदम सटीक टोकन-दर-टोकन मिलान प्राप्त हुआ—फिर भी हर जनरेट किया गया वाक्य निरर्थक (gibberish) था। "मिलान" (matching) और "काम करने" (working) के बीच का यह अंतर अब उन सभी के लिए एक चेतावनी है जो विशाल LLMs को Amazon के कस्टम इन्फरेंस चिप्स पर चलाने की कोशिश कर रहे हैं।

टोकन-दर-टोकन मिलान पर्याप्त क्यों नहीं है

डेवलपर ने Inferentia डिवाइस से प्रत्येक आउटपुट टोकन की तुलना मॉडल के CPU रन द्वारा उत्पादित टोकन से की। स्ट्रीम्स बिल्कुल समान थीं, इसलिए ऐसा लगा कि हार्डवेयर ने रेफरेंस इम्प्लीमेंटेशन को बिल्कुल सटीक रूप से पुनरुत्पादित किया है। वास्तव में, दोनों स्ट्रीम्स ने एक ऐसे मॉडल में गलत तरीके से बनाया गया प्रॉम्प्ट (malformed prompt) भेजा था, जिसे उसके चैट टेम्पलेट से हटा दिया गया था और गलत टर्न मार्कर्स दिए गए थे। टेम्पलेट की कमी ने मॉडल को एक अनंत लूप (infinite loop) में डाल दिया, जिससे वह निरर्थक बातें उगलने लगा। हार्डवेयर ने अपना काम किया—उसने उस बग को पुनरुत्पादित किया जो रेफरेंस कोड में पहले से मौजूद था।

सबक सरल है: SEQ_MATCH (सीक्वेंशियल टोकन समानता) का अर्थ शुद्धता (correctness) नहीं है। यदि रेफरेंस इम्प्लीमेंटेशन ही खराब है, तो एक वफादार हार्डवेयर रेप्लिका भी उसी विफलता को अपना लेगा। वैलिडेशन को टोकन-लेवल समानता से आगे बढ़ना चाहिए; इसके लिए सही ढंग से फॉर्मेट किए गए इनपुट्स के साथ एंड-टू-एंड फंक्शनल चेक की आवश्यकता है।

पैरामीटर्स के रूप में छिपते बफ़र्स

लोड चरण के दौरान, मॉडल लोडर ने layer_scalar नामक घटक को छोड़ दिया। कोड ने PyTorch मॉडल परिभाषा में इस ऑब्जेक्ट को पैरामीटर के बजाय एक बफ़र के रूप में रजिस्टर किया था। बफ़र्स स्टैटिक टेंसर होते हैं जिन्हें ट्रेनिंग के दौरान अपडेट नहीं किया जाता है, और कई लोडर्स Neuron-संगत फॉर्मेट में बदलते समय उन्हें अनदेखा कर देते हैं। इसे छोड़ने से कई लेयर्स के स्केलिंग कारक (scaling factors) उनके डिफॉल्ट मानों पर रह गए, जिससे पूरे नेटवर्क में गणितीय गणना बिगड़ गई। कोई एरर नहीं आया; मॉडल कंपाइल हुआ, और इन्फरेंस पाइपलाइन चली, लेकिन संख्यात्मक परिणाम (numerical results) गलत थे।

जो कोई भी बड़े मॉडल्स को Inferentia पर ले जा रहा है, वह हर नॉन-पैरामीटर टेंसर का ऑडिट करे। भले ही किसी टेंसर को सीखने (learn) के लिए न बनाया गया हो, फिर भी वह सही फॉरवर्ड-पास गणना के लिए आवश्यक हो सकता है। बफ़र के समावेश को मैन्युअल रूप से सत्यापित करने से उन साइलेंट स्केल एरर्स को रोका जा सकता है जिन्हें अन्यथा पहचानना कठिन होता है।

स्पॉट-इंस्टेंस की अस्थिरता और 39-मिनट का कंपाइल

स्पॉट इंस्टेंस पर 31-बिलियन-पैरामीटर वाला मॉडल चलाना सस्ता लगता है, लेकिन इस बचत के साथ अप्रत्याशित रिकलेम इवेंट्स (reclaim events) भी आते हैं। डेवलपर का कंपाइल समय—मॉडल को Neuron-संगत कोड में बदलने के लिए लगभग 39 मिनट—तब बर्बाद हो गया जब AWS ने इंस्टेंस को वापस ले लिया। रुकावटों से बचने के लिए उन्होंने एक तीन-स्तरीय सुरक्षा जाल बनाया:

  • ModelBuilder ने मेमोरी उपयोग को 384 GB होस्ट सीमा के भीतर रखा, जिससे उन क्रैश से बचा जा सके जो रीस्टार्ट करने के लिए मजबूर कर देते।
  • रॉ वेट फाइल्स (raw weight files) और कंपाइल किए गए “neffs” (Neuron executable files) दोनों का तत्काल S3 मिररिंग करने से एक नया इंस्टेंस ठीक वहीं से शुरू कर सका जहाँ पिछला वाला रुका था।
  • एक मल्टी-रीजन पोलर ने उपलब्ध स्पॉट क्षमता के लिए AWS रीजन्स को स्कैन किया और जैसे ही कोई उपलब्ध हुआ, एक नया इंस्टेंस लॉन्च कर दिया।

इन कदमों ने एक नाजुक, सिंगल-पॉइंट कंपाइल को एक लचीली पाइपलाइन में बदल दिया जो स्पॉट मार्केट के उतार-चढ़ाव में भी टिकी रहती है।

मिश्रित अटेंशन लेआउट के साथ शार्डिंग की कमियां

Gemma-4 31B दो अटेंशन कॉन्फ़िगरेशन का उपयोग करता है। कुछ लेयर्स चार की-वैल्यू (KV) हेड्स का उपयोग करती हैं, जबकि अन्य में संख्या अलग होती है। आठ समानांतर रैंकों (parallel ranks) में मॉडल को समान रूप से विभाजित करना तब विफल हो जाता है जब किसी लेयर के KV हेड की संख्या पूरी तरह से विभाजित नहीं होती है। 4-हेड वाली लेयर को आठ रैंकों में शार्ड करने की कोशिश करने से प्रत्येक रैंक को आधे हेड को संभालना पड़ता—जो कि एक गणितीय असंभवता है और इससे शेप मिसमैच और रनटाइम एरर उत्पन्न होते हैं।

समाधान यह था कि ग्लोबली-शार्ड किए गए लेयर्स (वे जिनकी हेड संख्या संगत है) को सभी रैंकों में रेप्लिकेट किया जाए और केवल उन “स्लाइडिंग” लेयर्स को शार्ड किया जाए जिनके हेड काउंट से समान विभाजन संभव हो। इस हाइब्रिड रणनीति ने KV हेड्स के अवैध विभाजन से बचते हुए टेंसर-पैरेलल दक्षता बनाए रखी, जिससे वे टेंसर-पैरेललाइजेशन एरर समाप्त हो गए जो पिछले प्रयासों में बाधा बन रहे थे।

निष्कर्ष

एक विशाल LLM को Inferentia पर पोर्ट करना केवल कंपाइल-एंड-रन का अभ्यास नहीं है। इसके लिए टोकन समानता से परे कठोर फंक्शनल टेस्टिंग, हर टेंसर—पैरामीटर या बफ़र—के सही ढंग से हैंडल किए जाने का सूक्ष्म सत्यापन, और एक ऐसी परिनियोजन (deployment) रणनीति की आवश्यकता होती है जो स्पॉट-इंस्टेंस रिकलेमेशन का अनुमान लगा सके। अंत में, शार्डिंग को मॉडल की आंतरिक अटेंशन ज्योमेट्री का सम्मान करना चाहिए; अन्यथा, पैरेललिज्म जो गति का वादा करता है, वह साइलेंट विफलता का स्रोत बन जाता है।