תקשורת מומחים משולבים יעילה וניידת

אנו מציגים ספריית קוד פתוח ניידת בעלת ביצועים גבוהים לתקשורת מומחים משולבים (MoE) המשיגה ביצועים מהירים פי 10 בהשוואה לפרימיטיבי תקשורת סטנדרטיים מסוג All-to-All.

היישום שלנו כולל מספר חדשנויות טכנולוגיות מרכזיות המספקות יעילות תקשורת MoE מעולה:

תקשורת ביוזמת GPU (IBGDA): תומכת בתקשורת ישירה מ-GPU ל-NIC, מה שמפחית משמעותית את ההשהיה על ידי עקיפת המעורבות של המעבד (CPU)

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

ביצועי צומת יחיד המהירים ביותר: השהיה נמוכה פי 2.5 בהשוואה ליישום המהיר ביותר בעבר בתצורות של צומת יחיד

ביצועי רב-צומת יעילים וניידים: היישום שלנו משיג מהירויות גבוהות פי 10 בהשוואה לתקשורת All-to-All סטנדרטית. למרות שהוא איטי בכ-פי 2 מיישומים מתמחים במיוחד, הגישה שלנו מציעה ניידות טובות יותר בין גרסאות NVSHMEM וסביבות רשת (NVLink, CX-7 ו-EFA)

הספרייה היא בקוד פתוח לחלוטין וזמינה בכתובת https://github.com/ppl-ai/pplx-kernels.

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

מבוא

מודלי מומחים משולבים (MoE), כגון DeepSeek R1 ו-Mixtral 8x7B, משפרים את המודלים הצפופים על ידי הגבלת מספר המשקולות המופעלות עבור כל טוקן. במקום להשיג ספירות פרמטרים גבוהות על ידי הגדלת מספר השכבות או הגדלת גודל השכבות הליניאריות בפרצפטרון רב-שכבתי (MLP) של כל שכבת מפענח, מודלי MoE מחליפים את ה-MLP המסורתי במספר מומחים ובנתב (router). לדוגמה, מתוך 671B הפרמטרים של DeepSeek R1, רק 37B מוכפלים בטוקן נתון במהלך ההסקה (inference). הדבר מפחית משמעותית את כמות החישוב הנדרשת לפענוח טוקן ומשפר את ההשהיה בהשוואה למודל צפוף.

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

כדי למזער את ההשהיה במערכות MoE מבוזרות, ניתן לנצל מקביליות על פני התקנים מרובים, אך הדבר יוצר אתגרי תקשורת. מודלים כגון Mixtral 8x7B או Llama-70B מתאימים ל-8 התקנים בצומת יחיד, ונהנים מחיבורי NVLink מהירים ובעלי השהיה נמוכה (עד 900Gbps). עם זאת, מודלים גדולים יותר דורשים פיצול (sharding) של מומחים על פני צמתים מרובים תוך שימוש ב-InfiniBand (המגיע לשיא של 400Gbps), מה שיוצר אתגרי השהיה נוספים.

הגדלה

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

פיצול ומקביליות

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

הגדלה

מקביליות מומחים (EP), המודגמת באיור הראשון, ממקבלת רק את חישוב המומחים. מומחים שונים מוקצים להתקנים שונים, המחזיקים במשקולות שלהם. לאחר הניתוב, טוקנים נשלחים להתקן המתאים, והתוצאות נאספות ומצטברות לאחר מכן. מורכבות הניתוב תלויה במידת המקביליות בחלקים האחרים של המודל: שכפול שכבות אחרות עשוי לבטל לחלוטין את הצורך בניתוב, מכיוון שכל דרגה (rank) יכולה לבחור את הטוקנים מטבלת ניתוב משוכפלת מקומית. עם זאת, אם רק דרגה אחת מריצה ניתוב, נדרש שידור (broadcast) כדי לשלוח טוקנים, אינדקסים ומשקולות למומחים המתאימים להם. לבסוף, איסוף-הכל (all-gather) או שידור All-to-All מסנכרנים את טוקני הפלט עם כל דרגה שממשיכה את ביצוע המודל. יישום כזה הוא פשוט יחסית מכיוון ש-torch כבר חושף את הפרימיטיבי הנדרשים, אם כי ניתן להפיק תועלת מסוימת על ידי מיזוג הצבירה עם פעולת האיסוף האוספת את הטוקנים מהמומחים.

הגדלה

מקביליות של מומחים בלבד אינה מתרחבת באופן אידיאלי, שכן צמתים באשכול עשויים להיות לא פעילים בזמן שהמודל מריץ שכבות שאינן של מומחים, כגון attention, norm ו-sampling. עם זאת, גם החישוב של שכבות אלה, ובעיקר attention, יכול ליהנות מ- Tensor Parallelism ‏(TP). רוב המודלים מסתמכים על multi-head attention, כלומר שניתן גם לחלק את הראשים של attention ואת ההטלות המתאימות שלהם של Q, K ו-V בין מכשירים, תוך שכפול או איסוף של הפרוסות בין שכבות שונות. אם attention מפוזר על פני כל המכשירים, all-gather יכול לסנכרן את ה-activations, ולאפשר לשכפול את ניתוב הנתונים, מה שמצריך פרימיטיבי סנכרון דומים למקרה של מקביליות מומחים לצורך מימוש יעיל. עם זאת, יש מגבלות למקביליות ברמה זו, שכן הפחתת מספר ראשים של attention מתחת לסף מסוים תניב תשואות פוחתות.

הגדלה

כדי לנצל בצורה הטובה ביותר את כל ההתקנים ולתמוך במידה גבוהה מאוד של מקביליות מומחים, של עד 128 או 256 GPUs, נדרשת מקביליות נתונים (DP). במסגרת סכמה זו, ההתקנים שעליהם מפוצל המודל מקובצים כדי לטפל בבקשות במקביל, תוך חישוב קשב ותחזוקת מטמוני KV המפוצלים על פני הקבוצה המקומית שלהם. מספר מופעים של קבוצות מקבילות אלו משיתפים פעולה בהערכת מומחים, כאשר כל אחת מארחת תת-קבוצה שונה של משקולות המומחים. בהתבסס על מספר ראשי הקשב, קבוצה עשויה בדרך כלל להשתרע עד לגודל של צומת שלם, שכן תקשורת תוך-צומתית מהירה יותר. לדוגמה, באיור שלמעלה, דרגת DP אחת מטפלת באופן עצמאי בשתי בקשות, A ו-B, ומטפלת בקשב, נורמליזציה וכל רישום אחר עבור הבקשות. דרגת ה-DP האחרת מעבדת בקשה נפרדת C. עם זאת, הצומת הראשון מארח מחצית מהמומחים, בעוד שהצומת השני מארח את המחצית השנייה, ולכן לאחר הניתוב טוקנים מ-A ו-B עשויים להישלח לצומת השני ולהפך. הדבר מוביל לבעיית תקשורת דלילה: כל התקן עשוי לשלוח מספר שונה של טוקנים לכל דרגת יעד אחרת. פרימיטיבי קיימים מ-torch, בעיקר all_to_all, אינם מתאימים במיוחד, שכן הם עשויים לדרוש צורה כלשהי של ריפוד (padding) או סנכרון GPU-ל-CPU ושידור מטא-דאטה. כדי ליישם תקשורת ביעילות, נדרשים גרעינים מותאמים אישית כדי לעבד את מידע הניתוב וליזום תקשורת מההתקנים ללא התערבות מעבד. לאחר הניתוב, גרעין משלוח חייב לשלוח טוקנים לדרגות שאליהן הם אותבו, בעוד ש בצד השילוב יש לאסוף את ההפעלות ששייכות לבקשות בקבוצת ה-DP הנוכחית. בנוסף, יש לאזן את העבודה בתוך קבוצת DP שעשויה לכלול מספר התקנים, על מנת למזער את השידור ולשכפל נכון את הטוקנים בכל הדרגות עבור השכבות העוקבות.

NVSHMEM

NVSHMEM הוא יישום OpenSHMEM הספציפי ל-NVIDIA, המספק מתקני תקשורת ניידים בין התקנים המפשטים את המורכבות של החומרה הבסיסית. ה-API יכול לבטא קריאות וכתיבות מ-GPU ל-GPU, הממופות לפרימיטיבי של שכבות העברה בודדות. הן NVLink והן RDMA נתמכים, מה שמעניק מידה של ניידות ומאפשר לגרעיני NVSHMEM לפעול בתוך צומת או על פני אשכול של צמתים מרובים. בגרעינים שלנו, אנו משתמשים בכתיבות זיכרון מרוחקות ובפעולות איתות (signaling) כדי להעביר נתונים ולסנכרן את ההתקנים המשתתפים.

פעולות NVSHMEM בנויות סביב המושג של זיכרון סימטרי: הן פועלות על מאגרים (buffers) שהוקצו בכל ההתקנים המשתתפים בתקשורת בין-התקנית. התקן מקומי יכול להשתמש בכתובת הנגזרת בתוך המאגר המקומי שלו לצד האינדקס של דרגת היעד כדי לציין את היעד של פעולה. האיור להלן ממחיש מושג זה: שני ה-GPUs מקצים מאגרים סיממטריים באותו גודל, השומרים על מצביעי src אליהם. הדרגה הראשונה רוצה לשלוח 3 מספרים שלמים לשנייה, ומקמת אותם בתחילת המאגר: nvshmem_int_put_nbi גוזר את כתובת ההתחלה מהמאגר המקומי, ומציין את התקן היעד. הדרגה השנייה גוזרת היסט (offset) מהמאגר שלה, ושולחת אלמנט אחד להתקן הראשון, עם היסט של אחד. בעוד שכתובות היעד חייבות להיות תמיד מאגרים סימטריים שהוקצו באמצעות nvshmem_alloc, מאגרי המקור יכולים להיות אזורים שרותיים של זיכרון ההתקן, בתנאי שהם נרשמו מראש ב-NVSHMEM.

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

הגדלה

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

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

nvshmem_uint64_wait_until משמשת בצד הקבלה של אות (signal), כדי לסקור דגל עד שהיעד המרוחק מעדכן אותו.

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

NVSHMEM משתמשת ב-GPUDirect RDMA כך שרשת ה-NIC תוכל לקרוא ולכתוב ישירות בזיכרון ה-GPU להחלפת נתונים מבלי לערב מעבדים. יתרה מכך, ברשתות ConnectX NICs, NVSHMEM תומכת בתקשורת ביוזמת GPU (הידועה גם בשם Infiniband GPUDirect Async, או IBGDA), המאפשרת ל-GPU לפרסם פעולות RDMA ישירות ל-NIC ולסקור את תור ההשלמה של ה-NIC ישירות. בפלטפורמות שאינן תומכות בתקשורת ביוזמת GPU, NVSHMEM מריצה שרשור "פרוקסי" במעבד המארח כדי ליזום תקשורת ולסקור השלמה בשם ה-GPU. תוכנית NVSHMEM ניידת ללא קשר לשאלה האם תקשורת ביוזמת GPU נתמכת או לא. עם זאת, תקשורת ביוזמת GPU מקצצת משמעותית את ההשהיה מכיוון שהיא עוקפת לחלוטין את העיקוף למעבד.

גרעינים ניידים

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

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

הגדלה

הצורה היחידה של סנכרון גלובלי, כפי שמוצג באיור למעלה, מיושמת בגרעין קבלת-השילוב (combine-receive). ברגע שנתונים נשלחים מהשליחה, דרגה חוזרת וממשיכה בביצוע כאשר היא מקבלת את כל הטוקנים שהיא צריכה מהדרגות האחרות ונכנסת לגרעיני שליחת-השילוב ברגע שהם מעובדים. מחסום (barrier) בגרעין קבלת-השילוב מבטיח שאף דרגה אינה יכולה לרוץ קדימה ולהתחיל בשליחת-הניתוב בזמן שכל דרגה אחרת עדיין ממתינה לקבלת נתונים, דבר שעלול לגרום לחפיפה הרסנית. הסנכרון מתבצע בגרעין השילוב לשם פשטות, וכאופטימיזציה ניתן ליישם את אותו מחסום בין קבלת הניתוב ושילוב השליחה, מה שמפחית עוד יותר את ההשהיה בכך שהוא מאפשר ל-NICs זמן רב יותר ליישב את ההעברות בזמן שה-GPU מריץ גרעינים אחרים.

שני הגרעינים מפוצלים על פני כל ה-SMs הזמינים של ההתקנים: בעוד שגרעיני שליחת ניתוב וקבלת שילוב חייבים למקבל על פני ספירת טוקנים מקסימלית לכל דרגה (max_m), גרעיני קבלת ניתוב ושליחת שילוב חייבים לטפל במספר כה רב של טוקנים מכל דרגה ומומחה מקומי, עבור מקסימום של max_m * (num_experts // EP) * (EP // TP).

ניתוב (Dispatch)

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

הגדלה

בחלק השולח, warp יחיד של 32 שרשורים (threads) בכל בלוק אחראי על קריאת מידע הניתוב וספירת מספר הטוקנים שכל מומחה יעד צפוי לקבל מהדרגה הזו, ושליחת הספירה פלוס אחד באמצעות nvshmem_signal_op. הספירה מוגדלת באחד, שכן המעבר מאפס ללא-אפס בצד המרוחק מסמל את קבלת הספירות מדרגת מקור. במקביל, שאר ה-warps משתפים פעולה כדי להעתיק טוקנים לזיכרון סימטרי על פני כל הבלוקים במקביל, תוך אריזת הפעלות, גורמי קנה מידה והאינדקס שלהם בדרגה המקומית לתוך נתח רציף. האינדקס נדרש על ידי השולח המשולב כדי לקבוע את הכתובת שאליה ייכתב הטוקן. לאחר מכן, לאחר הבטחת העתקת כל הנתונים באמצעות מחסום על פני קבוצות ה-warp, ה-warps פועלים שוב באופן עצמאי, כאשר כל אחד שולח את אותו מאגר למומחה אחר במקביל. הטוקנים נשלחים באמצעות nvshmemx_putmem_signal_nbi_warp, המגדיל גם באופן אטומי את ספירת הטוקנים שנשלחו מהדרגה המקומית בהתקן המרוחק. בתוך קבוצת DP, מכיוון שכל דרגה מחזיקה העתק של הטוקן לשליחה, הניתוב מאוזן באופן שווה, כאשר כל התקן שולח תת-קבוצה של הטוקנים.

בצד הקולט, כל הקרנלים מסתנכרנים תחילה עם ה-ranks השולחים על ידי המתנה לשליחה של סך כל ספירות הטוקנים. לאחר מכן, כולם ממתינים עד שספירות הטוקנים שנשלחו, שהוגדלו אטומית, יגיעו לספירות הכוללות, דבר המציין שכל המטענים מה-ranks המקוריים נשלחו גם הם, הודות לסמנטיקה של הקריאה putmem. הקרנלים מבצעים poll על הספירות באמצעות nvshmem_uint64_wait_until, ומקבילים את הפעולה על פני כל הבלוקים וה-threads. לאחר מכן, מחסום חוצה-בלוקים מבטיח שאף בלוק לא יקרא את ה-buffers אלא אם כן כל הנתונים התקבלו כראוי. מפוזרים על פני בלוקים ומסונכרנים באמצעות מונה שהוגדל אטומית, הטוקנים מועתקים מהזיכרון הסימטרי אל buffers רציפים רגילים, שהופכים לקלט עבור הקרנלים הבאים המממשים את המומחים. ה-rank המקור, אינדקס המומחה ואינדקס הטוקן נשמרים בנפרד, וממקמים בדיוק את המיקום שאליו קרנל ה-combine צריך לשלוח את ה-activations. אף על פי שטוקנים מתוך קבוצת DP נשלחים ממכשירים שונים, כולם מקובצים יחד כדי להישלח הלאה אל המומחה המתאים.

שילוב (Combine)

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

הגדלה

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

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

בדיקות השוואתיות (Benchmarks)

אנו מעריכים את הגרעינים שלנו על אשכול של 128 GPUs המשתרעים על פני 16 צמתים המחוברים באמצעות InfiniBand ו-CX-7 NICs. אנו מנסים הן GPUDirect Async (IBGDA) והן Reliable Connection (RC) עם פרוקסי CPU. אנו משווים אותם ליישום DeepSeek שעבר אופטימיזציה גבוהה, וכן לפרימיטיבי הצפופים המסופקים על ידי PyTorch (באמצעות NCCL) או NVSHMEM.

הגדלה

למרות שעל פעולת All-to-All הצפופה של NVSHMEM, הביצועים של IBRC ו-IBGDA דומים (6378 מיקרו-שניות לעומת 6180 מיקרו-שניות), IBGDA מהיר משמעותית עם הגרעינים הדלילים. בחיבור הניתוב והשילוב יחד, IBGDA צורך 902 מיקרו-שניות ואילו IBRC לוקח 3223 מיקרו-שניות - שיפור של פי 3.6 בהשהיה. בעוד ש-All-to-All מוגבל על ידי רוחב פס (bandwidth bound), הגרעינים הדלילים משדרים סדרי גודל פחות נתונים, ומוגבלים על ידי השהיה (latency bound). על ידי הפעלת העברות רשת ישירות מ-GPU, ללא צורך בפרוקסי CPU שיתאם בין ה-GPU לבין ה-NIC, השהיית קצה-לקצה מופחתת משמעותית.

בעוד שהגרעינים הניידים שלנו איטיים בכ-פי 2 מגרעיני DeepSeek שעברו אופטימיזציה גבוהה, הם משפרים את ההשהיה פי 10 בהשוואה לגרעינים הצפופים (902 מיקרו-שניות לעומת 9944 מיקרו-שניות).

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

הגדלה

בצומת יחיד (EP8), NVSHMEM משתמשת ב-NVLINK לצורך תחבורה, ומספקת השהיה נמוכה יותר ותפס (throughput) גבוה יותר מאשר רשתות בין-צמתים. הגרעינים הניידים שלנו מהירים פי 2.5 בערך מ-DeepEP בצומת יחיד (186 מיקרו-שניות לעומת 481 מיקרו-שניות).

שיפורים נוספים

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

סיכום

הצגנו ספרייה ניידת ובעלת ביצועים גבוהים לתקשורת MoE המשיגה ביצועים מהירים פי 10 בהשוואה לתקשורת All-to-All סטנדרטית תוך שמירה על תאימות בין תצורות חומרה מגוונות. בפריסות צומת יחיד עם NVLink, הפתרון שלנו מציג השהיה נמוכה פי 2.5 מיישומים קודמים.

הגישה שלנו מאזנת בין ביצועים לניידות באמצעות חדשנויות מרכזיות הכוללות תמיכה בתקשורת ביוזמת GPU, ארכיטקטורת גרעין מפוצלת המאפשרת חפיפת חישוב-תקשורת, וניתוב טוקנים יעיל תוך שימוש בפרימיטיבי NVSHMEM מינימליים. למרות שהיא איטית בכ-פי 2 מיישומים מיוחדים במיוחד בהגדרות רב-צומת, הספרייה שלנו מציעה גמישות מעולה בסביבות רשת שונות (NVLink, CX-7 ו-EFA).

ככל שמודלי MoE ממשיכים להתרחב, אסטרטגיות תקשורת יעילות כמו שלנו יהיו חשובות יותר ויותר לפריסה מעשית. היישום המלא שלנו בקוד פתוח זמין בכתובת https://github.com/ppl-ai/pplx-kernels.

מקורות

DeepEP

דוח טכני של DeepSeek-V3

NVIDIA NVSHMEM

ציטוט

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}