אימון מודלי AI בזול יותר: סנכרון משקלים דליל משנה את כלכלת ה-RL למודלי שפה
פיתוח חדש בספריית TRL מציע דרך יעילה בהרבה לסנכרן משקלים בין שרתי אימון לשרתי הסקה באימון RL אסינכרוני. במקום להעביר בכל צעד קבצי מודל ענקיים, המערכת שולחת רק את המשקלים שבאמת השתנו, ומקטינה דרמטית את עלויות התקשורת וההמתנה.
הבעיה הגדולה של אימון RL למודלי שפה
אחת מנקודות הכאב הפחות מדוברות באימון מודלי שפה גדולים באמצעות למידת חיזוק היא לא רק כמות ה-GPU, אלא תנועת המשקלים בין רכיבי המערכת. באימון RL אסינכרוני יש בדרך כלל הפרדה בין השרת שמעדכן את המודל לבין שרתי ההסקה שמייצרים דגימות, תגובות וסימולציות. לאחר כל צעד אופטימיזציה, שרתי ההסקה צריכים לקבל את גרסת המודל החדשה. במודל של 7 מיליארד פרמטרים מדובר בעשרות גיגה-בייטים, ובמודלי חזית של טריליון פרמטרים מדובר כבר בסדר גודל של טרה-בייט לכל סנכרון.
רוצה להישאר מעודכן ב-AI?
הירשם לדיוור השבועי שלנו וקבל עדכונים, המלצות על כלים, חדשות ודוחות מיוחדים
בפוסט טכני שפרסמה Hugging Face מוצגת גישה חדשה ב-TRL שמנסה לשנות את הנחת היסוד הזאת: לא צריך לשלוח את כל המודל, אלא רק את ההבדלים האמיתיים בין צעד אימון אחד לבא אחריו.
למה רוב המשקלים בכלל לא משתנים
הבסיס לרעיון הוא תכונה מספרית של חישובי bf16. באימון RL, קצבי הלמידה נמוכים יחסית, ולעיתים קרובות העדכון שמייצר האופטימייזר קטן מדי מכדי לשנות בפועל את הייצוג הבינארי של המשקל לאחר עיגול ל-bf16. במילים פשוטות, המודל עובר עדכון מתמטי, אבל עבור רוב הפרמטרים הביטים עצמם נשארים זהים.
מחקרים עדכניים כמו PULSE הראו שבין שני צעדי RL סמוכים, מעל 98% ולעיתים סביב 99% מהמשקלים נשארים זהים ברמת הביט. זו אינה דחיסה הסתברותית ואינה קירוב אגרסיבי, אלא ניצול של העובדה שרק חלק קטן מהטנסורים השתנה בפועל. המשמעות הכלכלית דרמטית: במקום להעביר מאות גיגה-בייטים או טרה-בייט, אפשר להעביר קובץ דלתא קטן בהרבה.
הדלתא עוברת דרך Bucket במקום דרך רשת ייעודית
היישום החדש ב-TRL משתמש ב-Hugging Face Buckets כמחסן אובייקטים משותף. שרת האימון מייצר קובץ safetensors דליל שמכיל אינדקסים של הערכים שהשתנו ואת הערכים החדשים שלהם. שרת vLLM מוריד את הקובץ, משחזר את העדכון ומטעין את המשקלים הרלוונטיים. אחת לכמה צעדים נשמר גם עוגן מלא, כלומר checkpoint מלא שממנו אפשר לשחזר שרשרת דלתות.
ההיבט החשוב כאן אינו רק טכני. זו ארכיטקטורה שמחליפה תלות בתשתיות רשת יקרות, כמו RDMA או קלאסטר אחוד, במודל מבוזר שמבוסס על object storage. שרת האימון יכול לרוץ בענן אחד, שרתי ההסקה בענן אחר, וסביבת הסימולציה בכלל ב-Space נפרד. כולם מתקשרים דרך אותו מאגר קבצים, ללא צורך בחיבור ישיר ביניהם.
מה זה פותח לתעשייה
עבור סטארט-אפים וצוותי מחקר, זה עשוי להיות שינוי משמעותי. עד היום אימון RL מתקדם למודלי שפה דרש בדרך כלל תשתית מרוכזת, יקרה ומורכבת לתפעול. אם סנכרון משקלים דליל יהפוך לסטנדרט, ניתן יהיה לבנות מערכי אימון מבוזרים וגמישים בהרבה, שבהם GPU פנוי במקום אחד משרת אימון, בעוד שרתי rollout נפרסים לפי צורך באזורים אחרים.
גם עבור חברות ענן ומפעילי מודלים זו בשורה חשובה. צוואר הבקבוק עובר מרוחב פס גולמי לניהול יעיל של גרסאות, מטמון, דלתות ושחזור. בעולם שבו אימון מודלים נעשה יותר ויותר מבוזר, הפחתה של פי עשרות או פי מאות בתעבורת המשקלים יכולה להשפיע ישירות על עלויות, זמני ניסוי וקצב פיתוח.
עדיין לא סוף הדרך
הפתרון הנוכחי עדיין דורש שמירת עותקי bf16 בזיכרון CPU כדי לזהות ולשחזר שינויים, והאינטגרציה עם vLLM צפויה להשתפר כאשר טעינת דלתות דלילות תיתמך ישירות על גבי GPU. גם מדיניות העוגנים עדיין בסיסית יחסית, וייתכן שבעתיד נראה עוגנים אדפטיביים לפי רמת drift מצטברת.
ובכל זאת, הכיוון ברור: אימון RL למודלי שפה גדולים לא חייב להיות כבול לקלאסטר יחיד ולסנכרון מלא של משקלים. אם הגישה הזאת תבשיל, היא עשויה להפוך את אימון ה-RL האסינכרוני ממותרות של מעבדות ענק לכלי נגיש בהרבה עבור האקוסיסטם הרחב של AI.