איך Amazon Search M5 חסך 30% עבור עלות אימון LLM באמצעות AWS Trainium | שירותי האינטרנט של אמזון

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

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

כמו ארגוני ML רבים, מאיצים משמשים במידה רבה כדי להאיץ אימון והסקת DL. כאשר AWS השיקה מאיצים ייעודיים עם המהדורה הראשונה של Afer Inferentia בשנת 2020, צוות M5 החל במהירות לנצל אותם כדי לפרוס ביעילות רבה יותר עומסי עבודה של ייצור, חיסכון בעלויות והפחתת זמן האחזור. בשנה שעברה השיקה AWS את זה AWS Trainium מאיצים, אשר מייעלים את הביצועים לכל עלות לפיתוח ובניית דגמי DL מהדור הבא. בפוסט זה, אנו דנים כיצד M5 הצליחה להפחית את העלות לאימון הדגמים שלהם ב-30%, ומשתפים כמה מהשיטות המומלצות שלמדנו לאורך הדרך.

מופעי טריניום

עם ההתקדמות באקסלרטורים ייעודיים, אמזון מספקת גם מאיצים מושכים בצורה של AWS Inferentia ו-Trainium. כפי שמרמזים שמם, שבבים אלה מותאמים כדי לעלות על הצרכים של מסקנות ועומסי אימון, בהתאמה. להכשרה בקנה מידה גדול של מודלים של יסודות שמגיעים למיליארדי פרמטרים בגודל, Trainium מופעי Trn1 ו- Trn1n הן בחירות אידיאליות בשל המאפיינים שלהן. מופעי Trn1 מופעלים על ידי הטכנולוגיה המתקדמת ביותר NeuronCore-v2, ויש להם כמות עצומה של מחשוב מאיץ וזיכרון. ניתן לבחור מופעי Trn1n גם עבור כמות גדולה יותר של רוחב פס רשת (1,600 Gbs), כך שהם מתאימים באופן אידיאלי לאימון ביצועים תוך מחשבה על אופטימיזציה של עלויות.

כדי להשתמש במאיצים, אתה צריך שכבת תוכנה שתתמוך בהם. עם שבבי Trn ו-Inf, ה AWS נוירון SDK פותח מאיצים ייעודיים של אמזון בעזרת PyTorch XLA. PyTorch XLA ממיר את המצב הלהוט של PyTorch ליישום מבוסס גרפים של מצב עצלן. לאחר מכן נעשה שימוש בגרפים אלה ונערכים נוספים לשימוש עם המאיץ. PyTorch Neuron (חלק מה- Neuron SDK) מאפשר למשתמשי PyTorch לאמן את הדגמים שלהם על Trainium NeuronCores עם כמה שורות קוד.

דגם ועומס עבודה

צוות M5 מאמן ופורס מודלים בסיסיים וייצוגים אוניברסליים כדי לסייע לצוותים שונים ברחבי אמזון להביא עונג ל Amazon.com לקוחות. מודל אחד כזה הוא מודל מקודד טקסט ואחריו פרצפטרון רב-שכבתי (MLP) עם אינטראקציות תכונות מפורשות או מרומזות המוגדרות על ידי ארכיטקטורת הרשת העצבית עם מאות מיליוני פרמטרים הניתנים לאימון. מודל זה מאומן על מיליארדי אסימונים, ומשמש ליצירת מיליוני הטמעות בהגדרת הסקת אצווה לא מקוונת. ההטמעות הללו הן תשומות לשירות אמזון מול לקוחות ברמה 1.

התשתית לשימוש בצינור הייצור אצווה AWS עם אסטרטגיות תור של נתח הוגן, תוך שימוש ב-EFA multi-node trn1.32xlarge cluster כמחשוב לאימון מודלים. מבחינה פונקציונלית, צינור הייצור מבצע אימון מודלים מצטבר, הערכה של מודל מאומן והסקת אצווה לא מקוונת על המודל המאומן, והכל באמצעות PyTorch כספריית ה-DL הבסיסית.

שערים

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

• איכות הדגם - איכות הדגמים שלנו משפיעה ישירות על חווית הלקוח. אנו דורשים שיהיה הבדל של פחות מ-0.1% באיכות הדגם בין GPU ל-Trainium.
• תפוקת אימון - אנו מאמנים באופן איטרטיבי את הדגמים שלנו מעת לעת כדי לספק את החוויה הרעננה ביותר ללקוחותינו. אנו דורשים כי התכנסות מודל חייבת להיות מושגת בתוך פרק זמן מוגדר מראש (כגון שבוע אחד) כדי לעמוד ב-SLAs הייצור שלנו.

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

תסריט הדרכה

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

• מיקום של xm.mark_step() - xm.mark_step() מרכיב ומפעיל את גרפי החישוב שנאספו בעצלתיים. קורא mark_step פעמים רבות מדי יובילו למספר גדול יותר של גרפים קטנים, בעוד שהפעלתו מעט מדי תוביל למספר גרפים קטנים, אך גדולים. בהתאם ליישום שלך, התפוקה והיישום של הכשרת המודל שלך ישתנו בהתאם למיקום שלך xm.mark_step(). היישום שלנו מציב אחד xm.mark_step() לאחר מסירה קדימה ואחורה, ואחד לאחר צעד האופטימיזציה.
• עטיפת מטעין נתונים עם מטעין התקנים מרובי עיבוד XLA - זהו שלב קריטי שניתן לפספס בקלות. מטעין התקני ריבוי עיבודים torch_xla.distributed.parallel_loader.MpDeviceLoader טוען נתוני אימון בכל מכשיר XLA עם אפשרויות לטעינה מראש ולחפיפה של טעינת נתונים עם ריצות מכשיר לשיפור התפוקה. גם מטעין המכשיר מפעיל xm.mark_step() ולכן הוא מסוגל לבנות גרפים לטעינת נתונים למכשיר מהמארח.

קומפילציה עבור Trainium

באופן מסורתי, מחזור הפיתוח של המודל עם GPUs כולל ביצוע שינויים בדגם או בסקריפט האימון והרצה ישירות על מכשיר ה-GPU. מאיצים כגון Trainium המשתמשים ב-XLA דורשים שלב נוסף לפני שניתן להפעיל אימון מודל על המאיץ. ניתן להפעיל גרפי חישוב XLA רק לאחר הידור. באופן כללי, ישנן שתי דרכים לבצע קומפילציה זו: Ahead of Time (AOT), שבו אתה עוקב וקומפילציה של כל הגרפים תחילה ולאחר מכן מפעיל אותם, או Just In Time (JIT), שבו מתקבצים אחר הגרפים, הידור ופועלים כפי שהם. נתקלים בהם. ה- Neuron SDK מספק את שני אלה מהקופסה. בדרך כלל, הידור AOT מבוצע תחילה. לאחר אוסף זה מופעלים גרפים. אם נתקלים בגרפים חדשים, זמן הריצה של Neuron מפעיל הידור JIT לפני הפעלתם. כדי לבצע קומפילציה של AOT, מספק ה- Neuron SDK neuron_parallel_compile, כלי קומפילציה המחלץ גרפים מהרצת ניסיון של סקריפט האימון ומבצע הידור מקביל של AOT.

היבט חשוב של הידור AOT הוא להבטיח שלא ייווצרו גרפי חישוב חדשים במהלך האימון. מקור אחד של גרפי חישוב חדשים (ולכן הידור מחדש) הוא צורות דינמיות של קבוצות האימון במהלך אימון המודל. מצאנו ששימוש בצורות סטטיות ואצוות בגודל קבוע מבטל את הידורי זמן האימון ומשפר מאוד את תפוקת האימון ללא כל השפעה על דיוק המודל. על ידי אכיפת אילוצים כאלה על אימון, ראינו שנדרשים רק 4-5 שלבים של אימון מודל, שלב אחד של אימות המודל ובדיקת נקודת ביקורת של המודל פעם אחת כדי להתחקות אחר כל הגרפים במהלך הידור AOT. חשוב לציין שה- Neuron SDK מתפתח כל הזמן, ובעתיד יתמוך גם בצורות דינמיות.

יתרה מזאת, הגרפים שנאספו מאוחסנים ב- Neuron Persistent Cache בדיסק או ב-an שירות אחסון פשוט של אמזון דלי (Amazon S3). זה שימושי במיוחד עבור עומסי עבודה של ייצור שבהם ארכיטקטורת המודל ותצורת ההדרכה אינם משתנים. לכן, התקורה של הקומפילציה נוצרת רק פעם אחת. השימוש במטמון הוא פשוט כמו הגדרת דגל סביבה:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

מהדר Neuron מספק גם שלושה אפשרויות אופטימיזציה ברמת המהדר (O1, O2, O3) לאיזון זמן הידור ותפוקת ריצת המודל. O1 מאפשר אופטימיזציות ליבה בגרף המחשוב וממזער את זמן ההידור, O3 מספק תפוקה משופרת של ריצת מודל במחיר של זמן הידור גבוה יותר, ו-O2 (אפשרות ברירת המחדל) מהווה איזון בין השניים. במקרה השימוש שלנו, השתמשנו באופטימיזציה של O1 וצפינו בהפחתה של 86% בזמן ההידור ללא שינוי במדדי דיוק המודל, תוך צפייה בירידה של כ-5-7% בתפוקה בהשוואה לאופטימיזציה של ברירת המחדל (O2). בהתאם למקרה השימוש, אתה יכול לבחור רמות שונות של אופטימיזציה.

לסיכום, השתמשנו בדגלים הבאים להידור:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

תאימות לנקודות ביקורת

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

מכיוון שדגם ה-GPU נשמר באמצעות כלי עזר סטנדרטיים לשמירת מודל PyTorch, הצלחנו להשתמש בכלי הטעינה של PyTorch Checkpoint כדי לטעון את דגם ה-GPU במכשירי Trainium.

לדוגמה, ב-GPU/CPU, אתה יכול לשמור את הדגם עם הקוד הבא:

torch.save(model.state_dict(), PATH)

ואז אתה טוען את הדגם בחזרה ל- Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

באופן דומה, אתה יכול לשמור את הדגם ב- Trainium עם הקוד הבא:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

וטען את הדגם בחזרה על GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

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

יציבות תפעולית

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

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

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

עם מנגנוני הגמישות הללו, הצלחנו להשיג 98.5% אחוזי הצלחה בעומסי העבודה שלנו ב-Trn1, בהשוואה לאחוזי ההצלחה הקיימים שלנו בצנרת ה-GPU.

תוצאות

כדי לאמת את הדיוק של הדגמים שלנו, אתחלנו שני דגמים מאותה נקודת ביקורת של GPU, ואימנו את האחד ב-Trainium והשני ב-GPU דומה. שני הדגמים אומנו עם אותם הפרמטרים של אימון. מערך הנתונים המשמש לחישוב מדדים הוא מערך נתונים של holdout, ואנו מעריכים את דיוק המודל במערך נתונים זה בכל N שלבים גלובליים. ציר ה-X הוא הצעד הגלובלי, וציר ה-Y הוא דיוק המודל. ראינו הבדל של פחות מ-0.1% בדיוק המודל בכל נקודה בגרף הבא.

יתר על כן, כדי להעריך את העלות-תועלת של אימון המודל, אנו מעדיפים להשוות את זמן שעון הקיר שנדרש כדי להגיע להתכנסות המודל. אנו מאמינים שזה מספק תצוגה מעשית יותר של חיסכון בעלויות בהשוואה למדדים כגון עלות לאסימון, FLOPS שהושג/דולר וגורמים אחרים. בהתחשב בזמן האימון של trn1.32xl וניתן להשוות ענן מחשוב אלסטי של אמזון (Amazon EC2), ראינו ש-Trainium מציעה עלות זולה עד 30% להתכנסות המודל.

סיכום

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

הודות לשותפות ושיתוף הפעולה שלנו עם צוות Annapurna Neuron, צוות Amazon Search M5 הצליח לחסוך עד 30% בעלויות על ידי מעבר ל-Trainium. הצוות מסוגל להשתמש ב-Trainium ולהשיג איכות מודל ושוויון תפוקה עם מאיצים דומים בשוק. יכולת פעולה הדדית של נקודות ביקורת ושינויי קוד מינימליים עם תמיכה ב-XLA אפשרו ל-M5 לבחור בין מספר מאיצים לעומסי העבודה שלהם. זה אפשר לצוות M5 לנצל את כוח המחשוב הגדול של Trainium, ולבנות פתרונות אגנוסטיים של מאיצים כדי לשמח את לקוחות Amazon.com. מנקודת מבט תפעולית, Trainium הוכח ככשיר לתמוך בשירותי שכבה 1 בקנה מידה של אמזון. צוות M5 ממשיך להעביר עומסי עבודה נוספים ל-Trainium כדי לספק את הדגמים הטובים ביותר עבור אמזון בעלויות הנמוכות ביותר.

לסיכום, צוות M5 הצליח לבצע אימון ML חסכוני בדרגת ייצור על ידי הוספת Trainium לצי המאיצים. אנו ממליצים לך להסתכל על Trainium ומכשירי Neuron אחרים כמו AWS Inferentia כדי לקצור את היתרונות של סיליקון אמזון המיועד למטרת עומסי עבודה של ML. התחל בקלות עם אחד ממספר ההדרכות הכוללות דגמים שונים, כמו Lama 2, זמין ב-Trainium.

על הכותבים

אבהיננדאן פאטני הוא מהנדס תוכנה בכיר ב-Amazon Search. הוא מתמקד בבניית מערכות וכלים לאימון למידה מבוזר ניתנת להרחבה והסקת מסקנות בזמן אמת.

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

ג'רי מניל הוא מהנדס תוכנה ב-Amazon Search. הוא עובד על שיפור היעילות, החוסן והמדרגיות של תשתית ההדרכה המבוזרת.

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

RJ הוא מהנדס בתוך אמזון. הוא בונה ומבצע אופטימיזציה של מערכות עבור מערכות מבוזרות להדרכה ועובד על אופטימיזציה של אימוץ מערכות כדי להפחית את השהיה עבור ML Inference. מחוץ לעבודה, הוא בוחן את השימוש ב- Generative AI לבניית מתכוני מזון.

• הפצת תוכן ויחסי ציבור מופעל על ידי SEO. קבל הגברה היום.
• PlatoData.Network Vertical Generative Ai. העצים את עצמך. גישה כאן.
• PlatoAiStream. Web3 Intelligence. הידע מוגבר. גישה כאן.
• PlatoESG. פחמן, קלינטק, אנרגיה, סביבה, שמש, ניהול פסולת. גישה כאן.
• PlatoHealth. מודיעין ביוטכנולוגיה וניסויים קליניים. גישה כאן.
• מקור: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/