ב-AI2 מצאו דרך לחלק שעות GPU בלי לקנות עוד שרתים

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

בפוסט שפורסם בבלוג של Hugging Face, אנשי צוות התשתית של AI2 מסבירים כיצד ניהול של אלפי מעבדי NVIDIA באשכולות של 88 עד 1,024 מאיצים הוביל אותם למבוי סתום מוכר.

כשהמשאב במחסור, המשתמשים לומדים לתחמן את המערכת.

כולם רוצים עדיפות עליונה

המערכת הקודמת של המכון התבססה על מתזמן עדיפויות סטטי, לצד אפשרות להגן על משימות מפני עצירה יזומה. התוצאה הייתה טרגדיה קלאסית של נחלת הכלל. חוקרים שרצו לוודא שתהיה להם מכונה זמינה לבדיקת קוד (דיבאג) פשוט שלחו עבודות סרק שתפסו כרטיסים שעות מראש. בנוסף, תוך זמן קצר נרשמה אינפלציית עדיפויות מוחלטת: 100% מהמשימות סומנו ברמת HIGH, מה שרוקן את מדרגות העדיפות מכל משמעות והרעיב עבודות רגילות לחלוטין.

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

תקציב זמן במקום בעלות

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

המתזמן החדש מייצר תחושה כאילו קיבלנו תוספת של 30% כוח מחשוב. במתזמן הישן, אם לא היינו מנצלים את מגבלת המכונות שלנו, המחשוב הזה פשוט הלך לאיבוד.כריס קלארק, חוקר ב-AI2

המערכת עברה למודל Fair-Share היררכי המבוסס על חלון היסטורי נע של שבעה ימים, לפי עקרונות שמזכירים אלגוריתמים ותיקים כמו SLURM’s Fair Tree או מתזמני YARN’s Fair Scheduler. ראשי פרויקטים ומנהלים מקבלים תקציב שעות GPU שבועי ומחלקים אותו כלפי מטה. המערכת מפרידה בין זמן מוקצה, שיורד מהתקציב ומעניק הגנה ראשונית, לבין ריצות מזדמנות ללא תקציב, שיכולות לרוץ על כרטיסים פנויים בחינם אך חשופות לעצירה מיידית בכל רגע.

חוזה עבודה במקום מונופול

החידוש המרכזי שפתר את התקיעות אינו האלגוריתם עצמו אלא חוזה התזמון. בעבר, משימת אימון מבוזר יכלה לשבת על שרתים במשך שבועות בלי אפשרות לגעת בה. במסגרת החוזה החדש, כל משתמש מחויב להצהיר על זמן ריצה מינימלי הדרוש למשימה כדי לייצר התקדמות משמעותית (עד תקרת מקסימום שנבדקה של שמונה שעות), או להגדיר אפס עבור זמן לא מוקצה. במהלך הזמן המינימלי המוגדר המשימה מוגנת מפני עצירה.

ברגע שחלון הזמן המינימלי מסתיים, המשימה רשאית להמשיך לרוץ כל עוד התקציב מאפשר זאת, אך היא עשויה להיעצר במידת הצורך: המתזמן יכול לבצע preemption, לשחרר את הכרטיסים לטובת משימות אחרות, ולהחזיר את העבודה המקורית לתור להמשך ריצה מאוחר יותר. המנגנון הזה חסך למכון 74% מהתערבויות התחזוקה הידניות של צוות ה-on-call, כיוון ששרתים פגומים יכלו לנקז משימות באופן אוטומטי מיד בתום החלון המינימלי שלהן.

התקורה שלא סיפרו עליה

הרעיון של AI2 נשמע מפתה, אבל הוא לא מגיע בחינם. כדי שחוזה תזמון יעבוד מבלי לזרוק לפח שעות של עבודה יקרה, כל אימון מבוזר בארגון חייב לתמוך באופן מוחלט בחידוש ריצה ושמירת נקודות ביקורת מהירה. אם שמירת checkpoint לוקחת 20 דקות ומעמיסה את כל רוחב הפס של מערכת הקבצים המשותפת, עצירות חוזרות ונשנות של עבודות של 1,024 מאיצים יחנוק את האחסון הארגוני.

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

מה אפשר ליישם כבר מחר

לצוותי MLOps שמנהלים אשכולות GPU משותפים יש מה ללמוד מהמהלך הזה. לא צריך לפתח מתזמן מאפס כדי לאמץ את הכללים הכלכליים של AI2. מנהלי תשתיות שעובדים עם Slurm יכולים להפעיל את מנגנוני ה-Fair-Share ההיררכיים הקיימים, ולהגדיר חלונות בדיקה שבועיים במקום חלוקת מכונות נוקשה. בסביבות Kubernetes מודרניות, כלים כמו Kueue מאפשרים לבנות היררכיות תקציב דומות.

הצעד הראשון והדחוף ביותר הוא לחייב את המפתחים לתמוך בחידוש ריצה אוטומטי כתנאי כניסה לאשכול. ברגע שכל משימה יודעת להיעצר ולהמשיך בלי לאבד חישוב, התמריץ לתפוס שרתים ללא צורך מתפוגג מעצמו. התוצאה בסימולציות של AI2, שבהן זמני ההמתנה למשימות דיבאג קצרות נחתכו מכשש שעות לחמש דקות בלבד, מראה שהבעיה היא לרוב לא גודל האשכול אלא המשמעת שלו.

אז לא, אתם כנראה לא צריכים עוד שורת שרתי Blackwell בשביל שהחוקרים יפסיקו להתלונן על התור.