KV कैश से परे: नेटिव पायटॉच के साथ AWS Trainium2 पर i14 के लैपलास 3B रैखिक ध्यान मॉडल का प्रशिक्षण और सेवा
aws ·
यह लेख बताता है कि i14 ने TorchNeuron Native के माध्यम से प्राकृतिक PyTorch का उपयोग करके अपने Laplace 3B लीनियर एटेंशन मॉडल को AWS Trainium2 पर कैसे प्रशिक्षित और सर्व किया। पारंपरिक ट्रान्सफॉर्मर्स के विपरीत, जो बढ़ते की-वैल्यू कैश पर निर्भर करते हैं, Laplace 3B गेटेड लीनियर एटेंशन (GLA) का उपयोग करता है, जो संदर्भ की लंबाई से स्वतंत्र रूप से प्रति अनुक्रम 29.5 MB का स्थिर पुनरावृत्ति राज्य बनाए रखता है। इस कार्यान्वयन में तीन प्रमुख इंजीनियरिंग अनुकूलन आवश्यक थे: मानक PyTorch ऑपरेटरों के साथ GLA कुंज को फिर से लिखना, Muon ऑप्टिमाइज़र के लिए वितरित संचार को अनुकूलित करना, और कंपाइलर के लिए बैकवर्ड पास को स्टाटिक रूप से अनरोल करना। एकल trn2.48xlarge इंसटेंस पर प्रशिक्षण ने लगभग 1,20,000 टोकन प्रति सेकंड की गति और स्थिर हानि अभिसरण प्राप्त की। इनफेरेंस प्रदर्शन ने डिकोड चरण को कंपाइल करके ऑपरेटर डिस्पैच ओवरहेड को कम करने के द्वारा 1,280 समवर्ती अनुक्रमों पर प्रति सेकंड 11,544 आउटपुट टोकन तक पहुंचाया। ये परिणाम दिखाते हैं कि कस्टम आर्किटेक्चर AWS हार्डवेयर पर उच्च दक्षता प्राप्त कर सकते हैं और लंबे संदर्भ वाले अनुप्रयोगों के लिए भविष्यवाणी योग्य मेमोरी उपयोग प्रदान कर सकते हैं।
aws