TorchNeuron בבטא ציבורית: ‏AWS מביאה PyTorch נייטיב ל-Trainium

AWS השיקה ב-public beta את TorchNeuron, ‏backend של PyTorch בקוד פתוח שמריץ קוד PyTorch באופן טבעי על מאיצי Trainium. עד עכשיו, כדי להעלות מודלים ל-Trainium היה צריך לעבור דרך תהליך הקומפילציה של Neuron SDK, ולעיתים קרובות גם להתאים את הקוד. ההבטחה של TorchNeuron היא שחוקרים יוכלו לנסות דברים על Trainium בלי לשכתב את הקוד לפריימוורק אחר.

לצד ההשקה, AWS פרסמה case study של i14, שאימנה והריצה את Laplace 3B, מודל linear attention, על Trainium2 עם PyTorch טבעי. ב-linear attention, ה-KV cache שגדל עם הזמן מוחלף ב-state רקורסיבי קבוע, במקרה הזה 29.5MB לכל רצף, לא משנה כמה ארוך ה-context. הבחירה בדוגמה לא שגרתית מכוונת: היא מראה ש-TorchNeuron מתמודד גם עם ארכיטקטורות חדשות, ולא רק עם מתכוני transformer מוכרים.

ההשקה היא חלק ממתקפה מתואמת של ענקיות הענן על החפיר שבנתה Nvidia סביב התוכנה שלה. Google חשפה באותו שבוע את TorchTPU ל-TPUs שלה, ו-Nvidia הגיבה בהדגשת התמיכה ב-PyTorch מהיום הראשון. מבחינת AWS, להקל על העבודה עם Trainium זה הכרחי לאסטרטגיית השבבים שלה, שעליה נשענות גם עסקאות אימון גדולות עם חברות ה-AI.

זו עדיין גרסת בטא. האימוץ יוכרע לפי ביצועים שמשתווים ל-CUDA בעומסים נפוצים, כיסוי ה-operators וכלי ה-debugging. כדאי לעקוב אחרי נתוני benchmark, אחרי מועד ה-general availability, ואחרי השאלה אם חברות ה-AI הגדולות שמתאמנות על Trainium יעבירו את קוד המחקר שלהן ל-TorchNeuron.