TRL v1.15 משיק ראש LM מאוחד שחוסך עד 82% זיכרון ומאפשר רצפים ארוכים פי 7

ספריית TRL של Hugging Face משחררת את גרסה 1.15 עם מה שמפתחי הפרויקט מגדירים כאופטימיזציה הגדולה בתולדותיה: ראש שפה מאוחד (Fused LM Head) שמופעל כברירת מחדל בכל מאמני ה-RL והזיקוק. במקום לבנות טנזור לוגיטים מלא במימדי [אצווה, רצף, אוצר מילים], גרעין Triton מקרין את המצבים הנסתרים דרך ראש השפה באריחים ומחשב ישירות הסתברויות לוגריתמיות ואנטרופיה פר-טוקן. התוצאה המעשית: ירידה של עד 82% בזיכרון שיא (peak VRAM) ואפשרות לאמן רצפים ארוכים פי שבעה על אותו חומרה.

איך זה עובד, גרעין Triton במקום טנזור מלא

השינוי חל על SFT, DPO, KTO, GRPO, RLOO ו-Distillation. הגרעין החדש מייתר שלושה נתיבי חישוב קודמים, הנתיב המלא של הלוגיטים, הנתיב המקוטע של Liger Kernel, ומנגנון _forward_redirection, ומחליף אותם במעבר יחיד שכותב רק את מה שהמאמן צריך: log-probs, אנטרופיה, ובאופן אופציונלי גם log_sum_sq_probs למשקלות WPO, mean_logits ל-chosen/rejected, ו-is_top1 לדיוק טוקן ממוצע ב-DPO. כל הפלטים הנוספים נבדקו מול הלוגיטים המלאים בבדיקות אוטומטיות. הדרישה היחידה: Triton על GPU עם CUDA, ROCm או XPU תחת לינוקס.

המספרים בבדיקות, Gemma-3-1B על 79 GiB

הבדיקות נערכו על B300 בודד עם מקצה PyTorch מוגבל ל-79 GiB (מקביל לזיכרון שמיש של H100 80GB), באמצעות google/gemma-3-1b-pt עם משקולות bf16 אקראיות, אוצר מילים של 262 אלף טוקנים, גודל אצווה 1, gradient checkpointing ו-SDPA. באורך 8,192 טוקנים נמדדה ירידה של 52% עד 82% בזיכרון שיא, ותאוצה של 2.3% עד 10.9% בזמן צעד. האורך המרבי הניתן לאימון קפץ ב-DPO מ-10 אלף ל-59 אלף טוקנים, וב-GRPO מ-29 אלף ל-115 אלף. הרווח הגדול ביותר מופיע במודלים קטנים עם אוצר מילים גדול; SFT כבר השתמשה ב-chunked_nll שנמנע מלוגיטים מלאים, ולכן נותרה ללא שינוי.

שינויים שוברים ומסלול הגירה

ארבעת מאמני ה-RL, GRPO, RLOO, DPO, KTO, מסירים את הנתיבים הישנים. מתאם PEFT על lm_head יזרוק שגיאה; הפתרון הוא modules_to_save=["lm_head"]. הפרמטר use_liger_kernel=True מוכרז כ-deprecated ויוסר ב-v2.0.0; הגרעינים של Liger עדיין חלים, אבל fused_linear_cross_entropy נכפה ל-False והגדרה מפורשת שלו תזרוק שגיאה, יש לעבור ל-model_init_kwargs={"use_kernels": True}. הפונקציות compute_metrics ו-compute_loss עם return_outputs=True עדיין מקבלות לוגיטים מלאים ממעבר קדמי נוסף, אך שתיהן מזהירות שזה יוסר ב-v2.0.0.

Checkpointing סלקטיבי ב-SFT

תוספת נפרדת בגרסה: כש-gradient checkpointing מופעל, פלט הקשב נשמר כעת במעבר הקדמי במקום לחושב מחדש במעבר האחורי. זה מחזיר את רוב ההאטה של checkpointing בהקשר ארוך, בעלות של טנזור נסתר אחד נוסף לכל שכבה. הגישה זהה ל-Eager SAC ש-torchtitan משתמשת בו תחת FSDP2, ואינה דורשת torch.compile.