גוגל ומטא מציגות את TorchTPU, מנוע PyTorch מקורי עבור TPU של גוגל

TorchTPU הוא שיתוף פעולה בקוד פתוח בין גוגל ומטא שמציג מנוע PyTorch מקורי עבור TPU של גוגל. הוא מאפשר הרצת מודלים והגירת TPU ישירות מהקופסה תוך שאיפה לביצועי הסקה ואימון תחרותיים בתעשייה באמצעות אופטימיזציות הרצה מיידית, אינטגרציית torch.compile, והאצת חומרה מקורית.

TPU של גוגל שימשו לאימון ופריסה של מודלי ML גדולים בקנה מידה עצום. בעוד TPU היו קשורים היסטורית ל-TensorFlow ו-JAX, מאמץ זה החל בגוגל באמצע 2025 עם צוות PyTorch במטא. האינטגרציה המוקדמת של PyTorch/XLA, שהוצגה ב-2020 והתבססה על טנזורים עצלים, דרשה שינויים פולשניים בקוד המקור, סבלה מחביון קומפילציה חמור, חסרה תמיכה בהרצה מיידית, ופיגרה בפרימיטיבים מבוזרים.

עקרונות עיצוב ושימושיות

TorchTPU נבנה כמנוע TPU עקרוני עבור PyTorch ולא כהתאמה של JAX. עקרונותיו כוללים אזרחות PyTorch (ללא תלות JAX בזמן ריצה, מצב מיידי מוצק, torch.compile המבוסס על Torch Dynamo, ותמיכה ב-DDP, FSDP ו-DTensor), אינטגרציה חלקה של המערכת האקולוגית, ומעורבות קהילה. המיקוד במשתמשים הוא פשטות העברה: שינוי model.to("cuda") ל-model.to("tpu") אמור לאפשר למודל לרוץ ישירות מהקופסה, אף שביצועי שיא עשויים לדרוש כוונון נוסף. הפרויקט גם חותר לדינמיות מוגבלת כדי להימנע מקומפילציה מחדש של XLA עקב שינויי צורה, וחוקר קרנלי TPU דינמיים שניתנים לקומפילציה מראש. ייבוא PyTorch על מכונת TPU כאשר חבילת TorchTPU מותקנת דורש רק שינויים זניחים דרך מנגנון הטעינה האוטומטית של PyTorch.

הרצה מיידית וקומפילציה

TorchTPU מרחיב את המפיץ של PyTorch במנוע C++ הרשום כ-PrivateUse1 ומגדיר אחסון טנזור TPU באמצעות מנגנון at::Tensor הסטנדרטי. הוא מציע שלושה מצבי הרצה מיידית: Strict (ברירת מחדל) שולח פעולת Aten אחת בכל פעם באופן אסינכרוני; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) הוא סינכרוני עם בדיקות גבולות נוספות; ו-DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) דוחה ומאחד מספר פעולות, מדווח שהוא מהיר ב-50% או יותר מ-Strict eager תוך הפחתת זיכרון HBM מרבי. בנקודת מטריאליזציה כגון קריאת .to("cpu"), TorchTPU מפעיל ניתוח חיות, בונה מפתח מטמון עבור גרף הפעולות, ושולח קובץ ביצוע שמור במטמון או מקמפל באופן אסינכרוני. היוריסטיקה של פעולות חוזרות מזהה רצפים חוזרים, כגון ערימות בלוקי טרנספורמר, ומפעילה מטריאליזציה מוקדמת. מטמון הקומפילציה התלת-שכבתי משתרע על זיכרון מקומי לתהליך, זיכרון משותף מקומי למכונה, ומטמון מערכת קבצים מרוחקת מתמשך לשימוש חוזר בין מארחים ואתחולים מחדש.

עומסי עבודה נתמכים ודיוק

TorchTPU תומך בכ-450 פעולות Aten, כולל מעברים קדימה ואחורה עבור פעולות Core ATen הנדרשות על ידי ארכיטקטורות מרכזיות, וכן קרנלים מותאמים אישית דרך JAX/Pallas או DSL של Helion. הוא אומת ישירות מהקופסה על 5,000 המודלים המורדים ביותר של Hugging Face, המכסים LLMs, אודיו ודיבור, ראייה, ומקודדי טקסט. שיתוף מבוסס DTensor, הסקה מרובת מארחים עם vLLM ו-SGLang, כוונון עדין ואימון עם TorchTitan, ואינטגרציה עם Ray ו-TorchRec אומתו גם הם. יחידות כפל מטריצה של TPU משתמשות בכניסות bfloat16 עם צבירת FP32, בעוד פעולות ללא כיווץ רצות ב-FP32 ביחידה הווקטורית. TorchTPU ממפה ממשקי API של דיוק PyTorch לחומרה זו: medium, ברירת המחדל, משתמש בכפל BF16 עם צבירת FP32; high משתמש באלגוריתם 3-מעברים בעלות חישוב של בערך פי 3; highest משתמש ב-6 מעברים בעומס של בערך פי 6. בחירת high או highest מפעילה אזהרת ביצועים.

מקורות: PyTorch · Google Cloud Blog · Hugging Face Blog