DeepSpeed עבור LLM: ZeRO, פריקת זיכרון, ריבוי GPU

הגדרת DeepSpeed: מדוע האימון נתקע בזיכרון

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

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

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

מה משתנה לאחר תצורת אימון נכונה

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

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

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

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

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

הגדרת DeepSpeed: מצבי ZeRO והעברת זיכרון החוצה

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

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

כיצד אנו בוחרים בפועל

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

סדר ההטמעה: מביקורת עומס ועד מסירה

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

שבעה שלבים

1. ביקורת האימון ופרופיל העומס. מצלמים פרופיל של צעד אחד: כמה זיכרון תופסים המשקולות, מצבי האופטימייזר והאקטיבציות, ואיזה חלק מהזמן המאיצים עומדים בטל בהמתנה לחליפין. בפלט — צוואר בקבוק שנקוב במספר, לא השערה.
2. סכימת הפריסה. מתאימים את דרגת החיתוך של מצבי המודל (ZeRO) ואת מצב ה-offloading: מה נשאר על המאיץ, מה עובר לזיכרון ה-CPU, ומה מחושב מחדש תוך כדי. את הסכימה מקבעים לפני התיקונים.
3. הטמעת האימון. משנים פרמטרים אחד אחד: גודל המיקרו-אצווה, צעד צבירת הגרדיאנט, חפיפה בין החליפין לחישובים, ותדירות נקודות הבדיקה. אחרי כל שינוי — מדידה.
4. סדרת ריצות ביקורת. ריצות קצרות על מערך נתונים אמיתי עם השוואת עקומת ההתכנסות לזו של הייחוס. המטרה — לוודא שהרווח במהירות לא נקנה בירידת איכות.
5. ייצוב ריצות ארוכות. מפעילים התאוששות מנקודות בדיקה, מגבילים את גידול מספרן, ובודקים את ההתנהגות בעת כשל של תהליך אחד — האימון ממשיך ולא נופל כולו.
6. תיעוד. סכימת הפריסה, ערכי הפרמטרים והנימוק לכל שינוי. בלי זה התצורה תחזור להיות «קופסה שחורה» אחרי התיקון הראשון.
7. מסירה לתפעול. מריצים את האימון יחד עם הצוות שלכם, מפרקים כשלים אופייניים ונקודות תיקון, ומשאירים נוהל הרצה ושינוי פרמטרים.

מה נשאר לצוות בסוף

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

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

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

מקרה בוחן: כיצד הסרנו את מגבלת הזיכרון באימון

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

מה מדדנו

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

מה שינינו ומה קיבלנו

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

את האקטיבציות אנו מחשבים מחדש במעבר לאחור ולא מחזיקים בזיכרון במלואן — החליפין הזה של חישוב בזיכרון משתלם כאן.

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

במה ZeRO שונה מהעברת מצבים לדיסק?

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

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

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

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

ספקות טכניים נפוצים לפני ההתחלה

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

האם זמן האימון יגדל?

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

כיצד לחזור אחורה אם התוצאה לא תספק?

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

כיצד לוודא שהאיכות לא ירדה?

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

נדון במשימה ונחזור עם הערכה ותוכנית

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

Для старта достаточно — להתחלה די ב:

• גישה למאגר הקוד עם סקריפטי ההרצה הנוכחיים ותצורות הפיצול וההעברה — די בהרשאת קריאה;
• תיאור המודל וגודלו, מספר המאיצים וסוג הקישור ביניהם;
• פרופיל עומס אמיתי: אורך הרצפים, שיעור ה-fine-tuning, דרישות לרוחב פס;
• יומנים של הרצות כושלות או דוחות על מחסור בזיכרון, אם נשמרו;
• מגבלות סביבה: תשתית עצמאית או ענן, דרישות לגישות ולבידוד;
• איש קשר של המהנדס שמקבל את התוצאה — כדי לתאם שינויים בלי מתווכים.

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

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

תארו את המשימה בצורה חופשית — נחזור עם הערכה ותוכנית עבודה.