אמזון הכריזה על TorchNeuron, קוד פתוח מקורי PyTorch למשפחת שבבי ה-AI של Trainium, בצעד שמסיר את אחד החסמים המעשיים הגדולים ביותר לאימוץ סיליקון מותאם אישית לעומסי עבודה של למידת מכונה. על פי פוסט רשמי בבלוג PyTorch שפורסם לקראת כנס PyTorch בצפון אמריקה ב-20 באוקטובר, מפתחים יכולים כעת להריץ קוד מודל PyTorch קיים ב-Trainium עם שינוי בשורה אחת, במקום לשכתב את זרימות העבודה שלהם סביב כלי עבודה קנייניים.

ההכרזה משכה תשומת לב מיידית מכיוון שהיא מכוונת למה שהאנליסטים תיארו זה מכבר כיתרון העמיד ביותר של Nvidia: לא ביצועי שבבים גולמיים, אלא מערכת האקולוגית של תוכנת CUDA הקושרת מפתחים למעבדי ה-GPU שלה. רויטרס מסגרה את המהדורה כמו אמזון ואמרה שבבי ה-AI שלה יעבדו עם מסגרת הפיתוח הפופולרית ביותר ב-AI, בעוד TechCrunch כינה את TorchNeuron ההזדמנות הגדולה ביותר של החברה ב-Nvidia עד כה. לקבלת ראייה רחבה יותר של האופן שבו שוק האקסלרטורים משתנה, סיקור חדשות החומרה של AI שלנו עוקב אחר ההתפתחויות הללו עם התפתחותן.

מה TorchNeuron משנה למעשה

ההבטחה המרכזית של TorchNeuron היא פשטות. בדוגמה שפורסמה בבלוג PyTorch, לולאת אימון קיימת פועלת על Trainium לאחר לא יותר מאשר העברת הדגם למכשיר החדש עם `.to('neuron')`. קוד שנכתב כנגד `torch.accelerator`, ה-API המומלץ של PyTorch למכשירים אגנוסטיים, פועל ללא שינוי, לפי הפוסט. עבור עומסי עבודה מהודרים, מפתחים מגדירים את ברירת המחדל האחורי פעם אחת עם `torch.compiler.set_default_backend("neuron")` ולאחר מכן מפעילים `model.compile()` ללא שינויים נוספים.

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

  • ביצוע להוט למחקר וניפוי באגים, עם "מצב להוט אדפטיבי" ניסיוני שמפחית את תקורה של שיגור
  • torch.compile לביצועי ייצור, לכידת מודלים לגרפים שהמהדר של Neuron מבצע אופטימיזציה, עם תמיכה מקורית ב-FlexAttention
  • ממשק ליבת הנוירון (NKI) לחוקרים הזקוקים לגרעין מותאמים אישית ברמה נמוכה

הקצה האחורי עובד גם על ה-Trainium2 וגם על השבבים החדשים יותר של Trainium3, על פי ההודעה.

פגיעה ישירה ב-Software Moat של Nvidia

האחיזה של Nvidia במחשוב AI מעולם לא נשענה רק על חומרה מהירה יותר. CUDA, פלטפורמת התוכנה בת שני עשורים של החברה, היא המקום שבו רוב קוד המחקר של AI נכתב, מכוון ונפרס, וכל שבב מאתגר - מקו האינסטינקט של AMD ועד ל-TPUs של גוגל ו-Trainium של אמזון - נאלץ להתמודד עם העלות של בקשת מפתחים לעזוב אותו. סיקור ההשקה של TorchNeuron ב-GeekWire, Benzinga ו-GuruFocus מסגר את ההכרזה באופן אחיד כהתקפה על החפיר הזה.

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

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

היקף טכני, ממחברות ועד אימון בקנה מידה אשכול

היכולות המפורטות בבלוג PyTorch מציעות שאמזון מכוונת הרבה מעבר לניסויים חובבים. TorchNeuron מכסה את מערכי המפעיל הנראים על פני מודלים של טקסט, מולטי-מודאלי ושפת ראייה, כמו גם יצירת תמונות ווידאו, מודלים של חלל מדינה ומודלים בעולם, עם דיוק BF16, FP32 ו-FP16; תמיכת TorchAO ו-FP8 נמצאת בגרסת בטא פרטית לפני הזמינות הכללית במהדורה הבאה.

אימון מבוזר, הדיסציפלינה שהופכת קוד GPU יחיד לריצות בקנה מידה של אשכולות, נתמכת באמצעות ממשקי ה-API הסטנדרטיים של PyTorch - DDP, FSDP2, SimpleFSDP, מקביליות טנזור וצינורות, וריסוק בסגנון ZeRO, יחד עם DTensor ו-DeviceMesh עבור קבוצות תהליכים רב-ממדיות. זרמים ואירועים ממפים את הסמנטיקה של PyTorch על תורי החומרה של Trainium, נכתב בפוסט.

להסיק, ערימות ה-Neuron של אמזון מתבססות על אותו בסיס: vLLM Neuron משרת מודלים של שפות גדולות, בעוד vLLM Omni Neuron משרת מודלים של דיפוזיה ליצירת תמונות ווידאו. ממשק ה-Neuron Kernel מאפשר לחוקרים לכתוב אופרטורים מותאמים אישית לביצוע נלהב וגם לביצוע קומפילציה, ו-AWS היא NKI-lib במקור פתוח, קבוצה של גרעינים מותאמים ל-Trainium המכסים חישובים בלוק של FlashAttention ו-MoE עבור תרחישים מיינסטרים LLM ורב-מודאליים.

כלי ביצועים כלולים גם כן. Neuron Explorer מרחיב את ה-`torch.profiler` הסטנדרטי כדי לעקוב אחר ביצוע ממסגרת לחומרה, ומקשר כל אירוע בחזרה לקו המקור שלו, ויכול לייצא מדדים בפורמט פרקט שניתן לצרוך את הערות הפוסט על ידי סוכנים לניתוח צוואר בקבוק אוטומטי.

זמינות ומה שיבוא אחר כך

TorchNeuron מתוכננת להשקת הקוד הפתוח שלה בכנס PyTorch North America ב-20 באוקטובר, לפי הפוסט בבלוג, עם תאימות שתרחיב עד Hugging Face Transformers v5. AWS מתאר את Trainium כפורטפוליו של שבבים שתוכננו מהיסוד לאימון והסקת AI בקנה מידה ענן.

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

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

---

הישאר לפני AI

קבל את החדשות, הניתוחים ופריצות הדרך האחרונות של AI - הכל במקום אחד.

קרא עוד חדשות AI →