העברת זיכרון GPU בביצועים גבוהים ב-AWS Sagemaker Hyperpod

תשתיות למידה עמוקה מודרניות דורשות לעיתים קרובות העברה של כמויות גדולות של נתונים בין מעבדים גרפיים (GPU) ברחבי מכונות. ב-Perplexity, נתקלנו באתגר טכנולוגי ייחודי: העברה יעילה של אזורי זיכרון GPU לא רציפים בין מכונות במהירות המרבית האפשרית. פלטפורמת היעד שלנו, מופעי AWS p5, מציעה רוחב פס רשת מרשים של 3200 Gbps דרך 32 כרטיסי רשת. מאמר זה חולק את מסע בניית פתרון הרשת המותאם אישית שלנו בעל הביצועים הגבוהים, המשיג 97.1% מרוחב פס תיאורטי זה.

האתגר הטכני

מקרה השימוש שלנו הציג מספר דרישות טכניות מרכזיות:

• העברה ברוחב פס גבוה בין מעבדים גרפיים מרוחקים של מקטעי זיכרון לא רציפים
• יכולת הוספה או הסרה דינמית של צמתים מפריסות Kubernetes מבלי לפגוע בפעולות השוטפות
• תמיכה בדפוסי תקשורת עמית-לעמית (peer-to-peer)

למרות שספריית NCCL של NVIDIA היא תקן ה-de facto ללמידה עמוקה מבוזרת, היא לא הייתה אידיאלית עבור מקרה השימוש שלנו:

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

רשתות מודרניות בעלות ביצועים גבוהים

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

רוב הרשתות שבהן אנו משתמשים מדי יום מבוססות על פרוטוקולי TCP/IP, שבהם יישומים תקשרים עם כרטיס הרשת דרך ליבת מערכת ההפעלה באמצעות שקעים (sockets). עם זאת, רשתות בעלות ביצועים גבוהים משתמשות ב-RDMA (גישה ישירה מרחוק לזיכרון) – מחסנית חומרה ותוכנה שונה לחלוטין המאפשרת גישה ישירה לזיכרון בין מכונות ללא מעורבות מעבד (CPU).

AWS מספקת את Elastic Fabric Adapter (EFA), ממשק רשת מותאם אישית המיישם את פרוטוקול המותאם של Amazon הנקרא Scalable Reliable Datagram (SRD). בניגוד לרשתות TCP/IP מסורתיות שבהן יש להעתיק נתונים פעמים רבות בין מרחב המשתמש, מרחב הליבה ומאגרי הרשת, EFA עם RDMA מאפשר העברת נתונים ישירה בין זיכרון ה-GPU לכרטיס הרשת, תוך עקיפת ה-CPU לחלוטין.

פילוסופיה של עיצוב רשתות בעלות ביצועים גבוהים

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

1. בעלות על מאגרים: בניגוד לשקעים מסורתיים שבהם הליבה מנהלת את מאגרי הרשת ודורשת העתקה בין מרחב המשתמש למרחב הליבה, RDMA דורש מיישומים לנהל את המאגרים של עצמם. כאשר יישום יוזם פעולת רשת, הוא מעביר את הבעלות על המאגר לכרטיס הרשת עד להשלמת הפעולה, מה שמבטל את הצורך בהעתקת נתונים.
2. רישום זיכרון: יישומים חייבים לרשום אזורי זיכרון מול ליבת מערכת ההפעלה. הליבה מגדירה מיפוי כתובות וירטואליות המאפשר ל-CPU, ל-GPUs ולכרטיסי הרשת להבין כולם את אותן כתובות וירטואליות. רישום זה הוא פעולת חד-פעמית המאפשרת העברות נתונים עוקבות ללא העתקה (zero-copy).
3. מישור בקרה לעומת מישור נתונים: רשתות בעלות ביצועים גבוהים מפרידות פעולות לשתי קטגוריות:
   1. פעולות מישור בקרה (כמו הגדרת חיבור ורישום זיכרון) עוברות דרך הליבה כדי להבטיח אבטחה
   2. פעולות מישור נתונים (העברת הנתונים בפועל) עוקפות את הליבה לקבלת ביצועים מרביים
4. קבלה לפני שידור: בהיעדר מאגרים המנוהלים על ידי הליבה, יישומים חייבים לפרסם מראש פעולות קבלה, תוך ציון היכן יש למקם את הנתונים הנכנסים. זהו שינוי מהותי ממודל השקעים שבו יישומים יכולים לקבל נתונים בכל עת.
5. השלמה מבוססת-סקר: במקום להמתין להתראות של הליבה באמצעות מנגנונים כמו epoll, יישומים מבצעים סקר ישיר על תורי ההשלמה של החומרה. כך מתבטלת התקורה של קריאות מערכת ומתקבלת תגובה מיידית לפעולות שהושלמו.
6. מודעות לטופולוגיית חומרה: הבנה ואופטימיזציה עבור טופולוגיית החומרה הן קריטיות להשגת ביצועים מרביים.

הבנת טופולוגיית החומרה

למופעי AWS p5 יש ארכיטקטורה פנימית מתוחכמת. כפי שמוצג להלן, כל מופע מכיל שני שקעי CPU היוצרים שני צמתים מסוג NUMA, כאשר כל צומת NUMA מתחבר לארבעה מתגים מסוג PCIe:

הגדלה

תחת כל מתג PCIe, אנו מוצאים:

• ארבעה כרטיסי רשת EFA בנפח 100 Gbps
• מעבד גרפי (GPU) אחד מסוג NVIDIA H100
• כונן NVMe SSD אחד

נתיבי הנתונים עבור העברות TCP/IP ו-RDMA מדגימים הבדלים בסיסיים בגישות שלהם:

הגדלה

בפרוטוקול TCP/IP (צד שמאל), יש להעתיק נתונים מספר פעמים:

1. מזיכרון ה-GPU למאגר היישום במרחב המשתמש
2. ממאגר מרחב המשתמש למאגר הליבה
3. ממאגר הליבה לכרטיס הרשת
4. תהליך הפוך מתרחש בצד המקבל

כל פעולת העתקה צורכת מחזורי CPU ורוחב פס של הזיכרון. היישום נדרש גם לבצע החלפת הקשר (context switch) בין מרחב המשתמש למרחב הליבה עבור כל פעולת רשת.

לעומת זאת, RDMA (צד ימין) מספק העברת נתונים אמיתית ללא העתקה (zero-copy):

1. כרטיס הרשת קורא ישירות מזיכרון ה-GPU
2. הנתונים עוברים ישירות לכרטיס הרשת המרוחק
3. כרטיס הרשת המרוחק כותב ישירות לזיכרון ה-GPU של היעד
4. היישום צריך רק לבדוק תור השלמה (CQ) במרחב המשתמש כדי לדעת מתי ההעברה הושלמה

עם RDMA ושיוך חומרה מתאים, העברת נתונים בין שני מעבדים גרפיים דורשת רק מעבר דרך מתג ה-PCIe המקומי והרשת:

הגדלה

בניגוד לכך, העברות TCP/IP חייבות להעתיק נתונים פעמים רבות דרך הזיכרון הראשי, מה הגורם לגודש משמעותי באפיק ה-PCIe:

הגדלה

בנייה באמצעות libfabric

השתמשנו ב-libfabric – מסגרת המספקת ממשק גנרי לשירותי fabric. היישום שלנו משתמש בשני סוגים של פעולות RDMA:

• RDMA דו-צדדי (SEND/RECV) עבור הודעות בקרה הנושאות מטא-דאטה לגבי אזורי זיכרון (למשל, היסטים וגדלים).
• RDMA WRITE חד-צדדי עבור העברת הנתונים בפועל, כאשר כל פעולת WRITE מטפלת במקטע זיכרון רציף אחד.

המסע לביצועי שיא

הפיתוח שלנו התקדם דרך מספר שלבים:

1. יישום העברת הודעות חד-כיוונית בסיסית באמצעות SEND/RECV
2. הרחבה לתקשורת דו-כיוונית
3. הוספת GPUDirect RDMA WRITE עבור העברה ישירה מ-GPU ל-GPU
4. הרחבה לטיפול במספר העברות מקבילות
5. הצגת תור פעולות לצורך חוסן
6. השגת ניצולת רוחב פס של 97.4% בכרטיס רשת בודד

בעת התאמה ל-32 כרטיסי רשת, יישמנו מספר אופטימיזציות קריטיות:

1. תורי פעולות: במקום להגיש פעולות ישירות לכרטיסי רשת, אנו מנהלים תור ברמת היישום. הדבר מספק חוסן נגד עומס ברשת ומפשט את מודל התכנות.
2. חימום רשת: יצירת חיבורים מראש שיפרה את ביצועי האתחול.
3. ריבוי שרשורים (Multi-threading): שרשורים ייעודיים עבור פעולות הרשת של כל GPU.
4. שיוך ליבות CPU: קשירת שרשורים לליבות CPU ספציפיות כדי להימנע מהשפעות NUMA ופספוס מטמון (cache misses).
5. שארדינג של מצב (State Sharding): הפחתת התחרות בין שרשורים.
6. צירוף פעולות למקבצים (Batching): הגשת פעולות מרובות יחד.
7. פרסום פעולות עצל (Lazy Operation Posting): פעולות מוכנסות תחילה לתור ביישום. לאחר סקר (polling) של תור ההשלמה, אנו מנסים להגיש פעולות ממתינות לכרטיס הרשת, מה שמבטיח שימוש יעיל במשאבי הרשת.
8. הקצאת מששאבים מודעת NUMA: הקצאת משאבי libfabric כמו תורי השלמה בצומת ה-NUMA הנכון כדי למזער את השהיית הגישה לזיכרון.

באמצעות אופטימיזציות אלו, השגנו ביצועים סופיים של 3,108 Gbps – שהם 97.1% מהמקסימום התיאורטי של רוחב הפס.

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

סיכום

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

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

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

ציטוט

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}