מערכת Tunix של Google מסירה את צוואר הבקבוק שמנע מלמידת חיזוק סוכנותית (agentic RL) בקנה מידה גדול להשתמש ב-TPUs ביעילות. על ידי הפרדת העבודה של יצירת נתוני אינטראקציה מהעבודה של עדכון המדיניות (policy), Tunix מרימה את ניצול ה-TPU מאחוזים בודדים לכיוון קיבולת מלאה כמעט, ובכך מפחיתה דרמטית את בזבוז המחשוב.
צוואר הבקבוק ב-agentic RL
agentic RL שונה מאימון מודלי שפה המוכר יותר בשיטת "next-token". סוכן (agent) חייב לשלוח קריאות API, להריץ קוד או להתקדם בסביבה מדמה, ולאחר מכן להגיב לתוצאה. לכן, לופ האימון הוא סינכרוני: המודל מייצר פעולה, הסביבה רצה, התוצאה חוזרת, ורק אז המודל מקבל עדכון גרדיאנט. כאשר שלב אחד בסביבה לוקח מספר שניות, חומרת ה-TPU היקרה נשארת במצב המתנה (idle), וניצול המשאבים המדווח יכול לצנוח מתחת ל-10%. חוסר היעילות הזה מתרגם ישירות לחשבונות ענן גבוהים יותר ולמחזורי מחקר איטיים יותר.
הארכיטקטורה המנותקת (decoupled) של Tunix
Tunix תוקפת את הבעיה על ידי פיצול שני השלבים — יצירת מסלולים (trajectory generation) ואופטימיזציה של המדיניות (policy optimization) — למאגרי חומרה נפרדים.
- שחקנים אסינכרוניים (Asynchronous actors) רצים על CPUs או GPUs זולים. כל שחקן מתקשר באופן רציף עם הסביבה שהוקצתה לו, מתעד פעולות ותצפיות, ומזרים את המסלולים (trajectories) שנוצרו לאחסון משותף.
- לומדים רציפים (Continuous learners) תופסים TPU Pods ייעודיים. הלומד שואב אצווה (batches) מהבאפר המרכזי ומבצע עדכוני גרדיאנט מבלי לחכות ששחקן מסוים יסיים rollout.
- באפר בעל תפוקה גבוהה (High-throughput buffer) נמצא באמצע, ומשמש כאזור הכנה למסלולים. מכיוון שהלומד יכול לקרוא נתונים במהירות שבה הבאפר מספק אותם, ה-TPU לעולם לא נתקע.
התוצאה נטו היא צינור אימון (training pipeline) שבו ה-TPUs נשארים עסוקים כמעט כל הזמן, ומקדמים את הניצול לעבר 100%.
מכשולים טכניים וכיצד Tunix מתגברת עליהם
פרקים (episodes) באורך משתנה וקומפילציה מחדש של XLA
הקומפיילר XLA של JAX מבצע אופטימיזציה עבור צורות טנזור (tensor shapes) קבועות. עם זאת, משימות agentic מייצרות רצפים באורכים שונים, מה שבדרך כלל היה מפעיל תהליכי קומפילציה מחדש יקרים. Tunix אורזת רצפים קצרים יחד ומקבצת פרקים בעלי אורך דומה לתוך "דליים" (buckets), מה ששומר על יציבות הצורות זמן מספיק רב כדי ש-XLA יוכל להשתמש מחדש בקרנלים (kernels) מקומפלים. התוצאה היא תפוקה יציבה ללא עומס הקומפילציה שהיה משתק את הביצועים.
הרחבה (Scaling) של מודלים מאסיביים על פני שבבי TPU רבים
אימון סוכנים עם יותר מ-70 מיליארד פרמטרים דורש פיזור של משקלים ונתונים על פני מספר צמתי TPU. Tunix משתמשת בפרימיטיב ShardMap של JAX כדי לפצל (shard) הן את פרמטרי המודל והן את האקטיבציות, מה שמאפשר ללומד להחזיק את המודל כולו בזיכרון ועדיין להזין לו נתונים במהירות גבוהה. אסטרטגיית פיצול זו מאפשרת לאמן מודלים שהיו בעבר מחוץ להישג ידם של TPU pod בודד.
גרדיאנטים מיושנים (Stale gradients) מצינורות מנותקים
כאשר השחקנים רצים לפני הלומד, הנתונים שהם מספקים עלולים להפוך ל"מיושנים" (stale) ביחס למדיניות הנוכחית. Tunix מפחיתה את הסטייה הזו באמצעות שני מנגנונים: דגימת חשיבות (importance-sampling) שנותנת משקל מחדש לדגימות ישנות כדי לשקף את הרלוונטיות שלהן, וסף מיושנות (staleness threshold) הניתן להגדרה שמשליך מסלולים שחורגים מגיל שנקבע מראש. יחד, הם שומרים על יציבות הלמידה גם כאשר הצינור פועל באופן אסינכרוני.
דברים שמאמצים צריכים לשים לב אליהם
- ביקורת שיהוי (Latency audit) – התועלת בניתוק (decoupling) תלויה בזמן התגובה של הסביבה. על צוותים למדוד שיהוי מקצה לקצה ולוודא שמאגרי השחקנים מותאמים בגודלם כדי לשמור על הבאפר מלא היטב.
- תכנון מאגר עובדים (Worker pool design) – CPUs או GPUs זולים יכולים לארח שחקנים רבים, אך הקצאת יתר (oversubscribing) עלולה לגרום לעומס על הרשת או על האחסון. מאגר מאוזן התואם את קצב הקליטה של הבאפר הוא חיוני.
- חוסן הבאפר (Buffer robustness) – האחסון המרכזי חייב להתמודד עם קצבי כתיבה וקריאה גבוהים מבלי להפוך לצוואר בקבוק חדש. בחירת מערכת אחסון עם שיהוי זנב (tail latency) נמוך ורוחב פס מספק היא חלק בלתי נפרד מהארכיטקטורה.
חסרונות פוטנציאליים
הארכיטקטורה המפוצלת מציגה יותר רכיבים נעים: ציי של חומרה נפרדת, באפר קבוע ולוגיקת תיאום כדי לאכוף מגבלות מיושנות.
שורה תחתונה
Tunix מראה שהעלות הדומיננטית ב-agentic RL אינה המודל עצמו, אלא זמן ההמתנה שנגרם על ידי לופים סינכרוניים של אינטראקציה. על ידי העברת עבודת ה-rollout לחומרה זולה והזנת TPU pod הלומד ברציפות מבאפר בעל תפוקה גבוהה, Google הפכה בעיה של ניצול נמוך מ-10% לתהליך עבודה בקיבולת מלאה כמעט.