TRL 1.13 משחרר מדריך לאימון הקשר של מיליון טוקנים על צומת יחיד

ספריית ה-RL הפתוחה TRL בגרסה 1.13 שמה במרכז את אימון ההקשר הארוך, עם מדריך מעשי שמדגים כיצד להריץ fine-tuning על רצפים של מיליון טוקנים ומעלה על צומת בודד של 8 GPU מדגם H100. השחרור כולל גם שיפורי מהירות וזיכרון כלליים, אבל הערך המיידי נמצא בפירוק הטכני של צוואר הבקבוק האמיתי בזיכרון, לא משקלי המודל, אלא חישוב ה-loss.

מה שמאחורי המספרים

הדוגמה במדריך מכוונת את Qwen3-8B על ספרים ממאגר PG-19 שמחוברים ברצף עד לאורך של כ-1.049 מיליון טוקנים לדוגמה. צעד אימון בודד לוקח כ-380 שניות (קצת יותר משש דקות), וה-loss ההתחלתי עומד על 4.31, ערך תקין שמעיד על קונפיגורציה נכונה; ריצה שגויה להקשר ארוך תתחיל סביב 10. בגישה נאיבית הרצף הזה דורש 288 GB לכל GPU; אחרי האופטימיזציות במדריך הצריכה יורדת ל-56 GB.

איפה הזיכרון באמת נגמר

הפרופיל מראה שהפיק בזיכרון מופיע בין המעבר קדימה לאחורה, בדיוק בנקודה שבה מחשבים את ה-loss. השכבה האחרונה של הדקודר מוציאה מצב נסתר בגודל (אורך_רצף × גודל_נסתר), שמוכפל במטריצת ה-language modeling head ויוצר מטריצת logits בגודל (אורך_רצף × גודל_אוצר_מילים), עשרות אלפי עמודות לכל טוקן. מטריצה כזו לא נכנסת בזיכרון כשהרצף מגיע למיליון טוקנים.

הפתרון: חישוב loss בחתיכות

במקום לממש את כל מטריצת ה-logits בבת אחת, TRL מחלק את החישוב למקטעים של 256 שורות בכל פעם. ה-cross entropy הוא סכום על שורות, וה-softmax רץ לאורך אוצר המילים של שורה בודדת, כך שכל מקטע מכיל את כל מה שנחוץ לחישוב החלק שלו ב-loss. הכפלה גדולה אחת הופכת לסדרת הכפלות קטנות, והפיק בזיכרון נעלם, ה-loss כבר אינו הצוואר בקבוק.

תנאים והסתייגויות

המדריך דורש transformers מגרסת main (לא גרסה משוחררת) כי הוא מסתמך על gradient checkpointing עם offload שעדיין לא נכלל בריליס רשמי. הקוד זמין להרצה מיידית עם accelerate והקונפיגורציה המצורפת, אבל מדובר בהדגמה טכנית, לא בפייפליין מוכן לייצור. אין עדיין מדדי ביצועים על משימות downstream אחרי אימון כזה, והמדריך עצמו מציין שהריצה נבדקה על צומת בודד בלבד.