AWS משיקה את TorchNeuron כדי להריץ PyTorch ישירות על Trainium

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

במקום להעביר מודל למסגרת נפרדת, מפתחים יכולים לבחור התקן Neuron מתוך PyTorch ולהמשיך לעבוד עם eager execution או עם torch.compile. TorchNeuron תומך ב-Trainium2 וב-Trainium3, וכולל רכיבים לאימון מבוזר ולחלוקת עומסים בין מאיצים. זהו שינוי בשכבת התוכנה; הוא אינו משנה את החומרה ואינו מבטיח שכל מודל ירוץ ללא תיקונים.

מה משתנה בקוד ובכלים

ה-backend נבנה כחיבור ישיר למנגנון ההתקנים של PyTorch. מבחינת צוות פיתוח, הקוד אמור להישאר קרוב יותר לגרסה שרצה על CPU או GPU, כולל אפשרות לבדוק פעולות במצב eager לפני קומפילציה. AWS מציגה זאת כדרך לקצר את המעבר בין ניסוי מקומי לבין אימון או הרצה על Trainium.

Hugging Face הודיעה ש-Transformers, Diffusers, Accelerate, PEFT ו-TRL מקבלים תמיכה ב-TorchNeuron. לפי הבדיקות שלה, 89 מתוך 100 ארכיטקטורות יצירת הטקסט הפופולריות ביותר שהיא בדקה נתמכות, לעומת 46 ב-Optimum Neuron. זהו נתון של Hugging Face על קבוצת מודלים מסוימת, ולא מדד כללי לכל עומסי PyTorch.

התמיכה ב-eager execution חשובה במיוחד באיתור תקלות. מפתח יכול להריץ פעולה, לקבל שגיאה במקום קרוב יותר למקור שלה ולבדוק אותה לפני קומפילציה מלאה. עבור צוותים שמחזיקים מסלול אימון אחד לכמה סוגי מאיצים, פחות קוד ייעודי עשוי להקל על תחזוקה ועל השוואות ביצועים.

גרסת הבטא עדיין חלקית

מסמכי ההשקה מפרטים מגבלות לא מעטות. במבחני התאימות של PyTorch יש כ-22 אלף מקרים שמסומנים כרגע לדילוג בגלל פעולות שעדיין אינן מורדות ל-Neuron. בין הפערים המוצהרים: tensors דלילים, צורות דינמיות וסמליות, forward-mode automatic differentiation וחלק מסוגי מסכות הקשב. גם במודל שנתמך, פעולה אחת שאינה ממומשת עלולה לדרוש שינוי בקוד או מעבר חלקי להתקן אחר.

מבחינה תפעולית, צוותים יצטרכו לבדוק שלושה דברים בנפרד: האם המודל מתקמפל, האם התוצאות נשארות בטווח הדיוק הנדרש, והאם העלות והקצב טובים יותר מהחלופה הקיימת. ההודעה אינה מספקת benchmark אחיד שמכסה אימון, inference, מודלים גדולים ועומסי מולטימדיה. ביצועים תלויים גם בגודל batch, באורך הרצף, בחלוקת המודל ובכמות התקשורת בין המאיצים.

המהלך נותן ל-AWS דרך ישירה יותר להכניס Trainium למסלול העבודה הרגיל של מפתחי PyTorch. הוא גם מאפשר לקהילה לראות את קוד ה-backend ולדווח על פערים. השאלות החשובות יתבררו לאחר שהבטא תהיה בשימוש: כמה מודלים באמת עוברים בלי התאמות, כמה זמן נחסך באיתור תקלות, ואילו עומסים מצדיקים מעבר מחומרה אחרת ל-Trainium.