גוגל ומטא משחררות את TorchTPU בקוד פתוח: PyTorch רץ על TPU בלי מעקפים

גוגל ומטא שחררו בקוד פתוח את TorchTPU, קצה עורפי (backend) מקורי של PyTorch עבור Cloud TPU, מאיצי ה-AI שגוגל מציעה בענן שלה. לפי ההכרזה, TorchTPU רץ במצב eager כברירת מחדל, מתקמפל דרך torch.compile ו-XLA, ומחליף את PyTorch/XLA בעומסי עבודה חדשים. ב-Hugging Face פרסמו באותו יום שהספריות המרכזיות שלהם כבר רצות עליו.

מה זה TPU, ולמה זה היה מסובך

TPU (Tensor Processing Unit) הוא שבב שגוגל מתכננת כבר כעשר שנים, ותפקידו אחד: חישובי המטריצות שרשתות נוירונים מבלות בהם את רוב הזמן. גוגל משתמשת בו כדי לאמן ולהריץ את המודלים שלה, ובהם Gemini. בענן אפשר לשכור משבב בודד ועד pods של אלפי שבבים.

הבעיה הייתה תמיד התוכנה. כפי שכותבים ב-Hugging Face, סביבת ה-TPU צמחה סביב TensorFlow ו-JAX, ומשתמשי PyTorch, אחת מספריות למידת המכונה הנפוצות בעולם, נאלצו להתאמץ יותר. עד היום הדרך להריץ PyTorch על TPU עברה דרך PyTorch/XLA (torch_xla). זה עבד, אבל עם מגבלות: הטנזורים היו "עצלים", כלומר הקוד לא תמיד רץ ברגע שהמפתח ציפה. שינויים קטנים יכלו לגרום לקימפול מחדש של גרף החישוב, והדיבוג היה קשה. מי שלא רצה להתמודד עם זה נאלץ לרוב ללמוד JAX.

מה משתנה

עם TorchTPU, ה-TPU הופך ל"עוד התקן" בתוך PyTorch, בדומה לכרטיס גרפי. כותבים קוד PyTorch רגיל, ומצב eager פירושו שכל פעולה רצה מיד, כמו שרוב המפתחים רגילים. כשרוצים יותר מהירות מפעילים את torch.compile, שמקמפל את הקוד דרך XLA. לפי Hugging Face, החבילה כוללת גם kernels של Pallas שמכווננים ל-TPU ותמיכה באימון מבוזר. ההתקנה על מכונה וירטואלית עם TPU היא פקודת pip אחת.

האקוסיסטם של Hugging Face

ב-Hugging Face מסבירים ש-Transformers v5 נבנתה כמקור הסמכות להגדרות מודלים, כדי שאפשר יהיה להריץ אותם בכל מקום, ועכשיו גם על TPU. כדי לטעון מודל, למשל Qwen3.5-9B, מספיק להעביר את הפרמטר device_map="tpu". ב-Diffusers, ספריית מודלי התמונה והווידאו, מספיקה הפקודה pipe.to("tpu"). להרצה בייצור מוסיפים קריאת compile אחת, והקריאה הראשונה מקמפלת גרף שנשמר לקריאות הבאות באותו גודל תמונה. מודלים שגדולים מדי לשבב אחד אפשר לפצל בין כמה שבבים כבר בזמן הטעינה.

החברה גם מודיעה ש-Optimum TPU, הכלי שפיתחה כדי לגשר על הפער, כבר לא נדרש להרצת מודלים ב-Transformers וב-Diffusers.

מבחני הדוגמה שפורסמו רצו על Cloud TPU v6e-4, עם המודל מפוצל בין ארבעה שבבים ובדיוק bf16. יצירת תמונה ברזולוציה 1024×1024 ב-40 צעדים עם Qwen-Image-2.1 ארכה 3.5 שניות במצב מקומפל, לעומת 8.0 שניות במצב eager. יצירת סרטון של 5.2 שניות עם אודיו ב-MiniMax-H3, ברזולוציה 1344×768, ארכה 188 שניות מקומפל לעומת 262 שניות ב-eager.

למה זה חשוב

המשמעות המעשית היא ככל הנראה הורדה של מחסום כניסה: מפתחים שעובדים ב-PyTorch יוכלו לנסות TPU בלי לשכתב קוד ובלי ללמוד מסגרת עבודה חדשה. ההשתתפות של מטא, שמובילה את פיתוח PyTorch, מרמזת כנראה שמדובר בתמיכה רשמית ולא בתוסף צדדי. מבחינת גוגל, ההערכה היא שהצעד נועד להפוך את ה-TPU לחלופה נגישה יותר לכרטיסים של אנבידיה, שרוב קוד ה-PyTorch נכתב עבורם. עם זאת, המספרים שפורסמו הם דוגמאות של Hugging Face עצמה ולא השוואה בלתי תלויה מול חומרה מתחרה.