העברת משקלים עבור פוסט-אימון RL בפחות מ-2 שניות

לאחרונה השגנו עדכוני פרמטרים חוצי-מכונות תוך 1.3 שניות עבור Kimi-K2 ‏(1T פרמטרים), עם העברת משקולות מ-256 GPUs לאימון (BF16) אל 128 GPUs להסקה (FP8).

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

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

הפתרון שלנו מבוסס על RDMA WRITE, פרימיטיב חד-צדדי שבו מקור הכתיבה כותב ישירות לזיכרון ה-GPU של היעד.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

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

איסוף מטא-דאטה – הבקר אוסף מטא-דאטה של פרמטרים מכל מעבדי ה-GPU של האימון וההסקה.

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

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

ביצוע – לאחר כל שלב אימון, הבקר מאותת למעבדי ה-GPU של האימון להתחיל בהעברות.

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

פרמטרים באימון מופצים בהתאם למיקומי FSDP. באמצעות full_tensor(), כל מעבדי ה-GPU ב-DeviceMesh יכולים לשחזור את הפרמטר המלא, ולכן כולם יכולים לשמש כמקור להעברת משקלים.

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

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

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

העתקת זיכרון מ-Host ל-device‏ (memcpy) — אם FSDP מעביר משקלים ל-CPU

הכנת פרמטרים — שחזור משקל מלא באמצעות full_tensor(), החלת מיזוג הטלה (projection fusion), וקימות (quantization) במידת הצורך.

העברת RDMA — כתיבה ללא העתקה לזיכרון ה-GPU של ההסקה המרוחקת

מחסום גלובלי — לאחר שכל קריאות ה-full_tensor() הסתיימו, סנכרון בין קבוצות mesh באמצעות GLOO על גבי Ethernet.

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

full_tensor() ופעולות GPU אחרות מייצרות שימוש נוסף בזיכרון ה-GPU. כדי להימנע משגיאת חוסר זיכרון (OOM), אנו מתחילים בביצוע משימה רק אם המשימות הנוכחיות המתבצעות תופסות פחות זיכרון GPU זמני מסף מוגדר (watermark).

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

העברת משקלים מבוססת Rank0 אוספת את המשקלים לפני הפצתם לצומת האינפרנס.

תקשורת נקודה-לנקודה מעבירה משקלים משתי צמתי האימון אל ממשקי הרשת של צומת האינפרנס.

דפוס נפוץ הוא לנתב את כל הפרמטרים דרך מעבדי GPU בדירוג 0: איסוף בדרגה 0 של האימון, שליחה לדרגה 0 של ההסקה, ואז פיזור מחדש. זה הופך במהירות לצוואר בקבוק, המוגבל על ידי רוחב הפס של PCIe וכרטיס הרשת (NIC) של GPU יחיד (לדוגמה, 400 Gbps ≈ 50 GB/s).

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

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

מערכות מסוימות מתבססות על קריאה למתודת update_weight() של מנוע ההסקה עבור כל טנזור. המשמעות היא שינויים פולשניים בקוד ההסקה, בתוספת תקורה מ-RPC, סריאליזציה ותיאום מישור הבקרה.

באמצעות פרימיטיב RDMA WRITE, אנו מעדכנים משקלים בשקט בזיכרון ה-GPU של ההסקה, ללא העתקות מיותרות. אין הודעות במישור בקרה ואין לוגיקת בקקרת CPU. אין צורך בשינוי של מנוע ההסקה.

תהליך העברת המשקלים יכול לנצל ארבעה סוגים של משאבי חומרה: (1) העברת נתונים בין מארח להתקן (Host-device) (2) חישובי GPU למיזוג הטלה וקימות (3) רשת RDMA למישור הנתונים (4) Ethernet למישור הבקרה.

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

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

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

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

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

העברת משקלים מהירה ואמינה היא לבנת בניין קריטית לכוונון עדין של RL בקנה מידה רחב. על ידי שילוב של פרימיטיב RDMA WRITE, לוח זמנים סטטי להעברה ו-ביצוע בצינור, צמצמנו עדכוני טריליון פרמטרים ל-1.3 שניות בלבד ב-Kimi-K2. הגישה קלה להבנה, פשוטה לתחזוקה, ומונעת את צווארי הבקצר של עיצובים מסורתיים.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: העברות חד-צדדיות

תהליך עבודה ברמה גבוהה

ביצוע העברת משקלים

DeviceMesh וקבוצות Mesh

צינור משימות

בקרת שימוש בזיכרון ה-GPU

למה זה מהיר ופשוט

תקשורת נקודה-לנקודה

העברת נתונים חד-צדדית

צינור עיבוד (Pipelining)

לוח זמנים סטטי

הפרדה נקייה

סיכום

ציטוט