אמן מודלים ענקיים עם קנה מידה כמעט ליניארי באמצעות מקביליות נתונים מרוסקים ב- Amazon SageMaker

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

זו הסיבה שהשקנו לאחרונה sמקביליות נתונים קשוחים on אמזון SageMaker, טכניקת אימון מבוזרת חדשה שחוסכת זיכרון ב- ספריית מודל מקבילית של SageMaker (SMP).. מקביליות נתונים מפוצלים נבנתה במיוחד עבור מודלים בקנה מידה קיצוני ומשתמשת באמזון MiCS טכנולוגיה מתחת למכסה המנוע, מאמץ מדעי למזער את קנה המידה של התקשורת על ידי הפחתת תקורה יקרה של תקשורת הנטועה באיסוף פרמטרים וסנכרון גרדיאנט. עם דגם GPT-30 פרמטר 2B עם אורך רצף 2048, תכונה חדשה זו השיגה 141 TFLOPs, מהירות של 39.7% בהשוואה ל-DeepSpeed ​​ZeRO-3. עבור דגם 10B GPT-2 עם אורך רצף 512, תכונה חדשה זו השיגה גם 564 דגימות בשנייה, מהירות של 13.9% בהשוואה ל-Fully Sharded Data Parallel (FSDP) של PyTorch. זכור שבאימוני דגמים ענקיים, כל אחוז מהאצה מתורגם לדולרים שנחסכו ולפרודוקטיביות שהושגה בצוות שלך.

בפוסט זה בבלוג, תחילה נסקור מקרוב את המבדילים העיקריים של מקביליות נתונים מרוסקים ומתי להשתמש בהם. לאחר מכן, תלמד כיצד להכשיר דגם GPT-30 פרמטר 2B ב- SageMaker בקלות עם התכונה החדשה הזו. לבסוף נשווה את הביצועים לאפשרויות אחרות של קוד פתוח, במיוחד עם ביצועים טובים יותר מ-DeepSpeed ​​ZeRO של עד 39.7% ב-256 GPUs.

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

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

במסגרות קיימות שנופלות תחת פרדיגמה זו, בעיקר DeepSpeed ​​ZeRO-3 וה-FSDP של PyTorch המועלה מ-FairScale, מצבי המודל מחולקים על פני את כל GPUs, אסטרטגיה שמקטינה את צריכת הזיכרון בכל GPU במחיר של תקורה של תקשורת גדולה אשר גדלה עם גודל האשכול ולכן גורמת לירידה משמעותית בקנה מידה בקנה מידה. לעומת זאת, מקביליות נתונים מרוסקים במודל המחיצות של ספריית SMP מצבים ב-a מודע לקנה מידה באופן על ידי חלוקת כל העתק של מצבי מודל רק בתוך תת-קבוצה של GPUs.

בואו נסתכל מקרוב על חלוקת מודלים עם מודעות לקנה מידה ב-MiCS, טכנולוגיית הליבה מאחורי נתונים מרוסקים במקביל. האינטואיציה מאחורי העיצוב הזה היא שייתכן שלא תידרש חלוקת מצבי אימון על פני כל קבוצת הנתונים המקבילים כדי לאמן מודל עם עשרות מיליארדי פרמטרים. לדוגמה, 8 מעבדי V100 (32GB כל אחד) מספיקים כדי להחזיק את מצבי הדגם העתק של דגם של 10B פרמטרים שזקוק לכ-200GB של זיכרון בעת ​​אימון עם Adam Optimizer תוך שימוש ברמת דיוק מעורבת. על ידי הגבלת העתק שלם של מצבי מודל ב הקטן ביותר תת-קבוצה של GPUs, נוכל להפחית ביעילות את היקף התקשורת תקורה בהשוואה ל-DeepSpeed ​​ו-PyTorch FSDP. נתונים מפוצלים מקבילים ממנפים גם טכניקות אחרות ב-MiCS כמו תקשורת היררכית וסנכרון 2-Hop Gradient. למידע נוסף, בדוק קנה מידה כמעט ליניארי של אימון מודל ענק ב-AWS or MiCS: קנה מידה כמעט ליניארי לאימון מודל ענק בענן ציבורי.

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

שימוש במקביל לנתונים מרוסקים לאימון GPT-2 ב-Amazon SageMaker

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

1. התחל בעבודה

מקביליות נתונים מפוצלים זמינה ב-PyTorch v1.12.0+ ועובדת גם עם FP16 וגם עם BF16. הדרך הקלה ביותר להשתמש בספריית SMP היא באמצעות מיכל AWS Deep Learning בנוי מראש עבור PyTorch. עם זאת, אם אתה רוצה להביא מיכל Docker משלך, אתה יכול לפנות אליו צור מיכל Docker משלך עם ספריית המודל המקבילית של SageMaker. כדי להתחיל, עקוב שנה סקריפט אימון של PyTorch כדי להתאים את ממשקי ה-API של SMPs בסקריפט ההדרכה שלך. בחלק זה, אנו קוראים רק כמה שלבים עיקריים עם קטעי קוד מסקריפט ההדרכה המוכן לשימוש train_gpt_simple.py. אתה יכול לעקוב אחר ההערות בתסריט ו מסמך API למידע נוסף על היכן משתמשים בממשקי API של SMP.

ראשית, ייבא ואתחל את הספרייה על ידי קריאה smdistributed.modelparallel.torch.init() בתחילת תסריט האימון:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

שנית, עטפו את הדגם שיש לחלק איתו smdistributed.modelparallel.torch.DistributedModel ולהשתמש במוחזר DistributedModel אובייקט קדימה:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

עטפו את האופטימיזציה עם smdistributed.modelparallel.torch.DistributedOptimizer לשמירה וטעינה של מצבי אופטימיזציה.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

שים את ההיגיון קדימה ואחורה בפונקציית צעד וקשט אותו עם smdistributed.modelparallel.torch.step.  כל חישוב המוגדר בתוך smp.step-decorated הפונקציה מבוצעת בצורה מבוזרת.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. הכן את מערך הנתונים

אנו משתמשים openwebtext הוא מערך הנתונים בו אנו משתמשים בדוגמה זו. המחברת משתמשת בתסריט data_prep_512.py כדי להוריד ולעבד מראש את מערך הנתונים. אתה יכול גם להתאמן עם מערכי נתונים אחרים על ידי שינוי data_pipeline.py. כאשר אתה מתמודד עם מערך נתונים ומודלים גדולים, אתה יכול להאיץ את עבודת ההדרכה על ידי שימוש בנתונים המאוחסנים ב אמזון FSx עבור ברק, המספק מערכת קבצים בעלת ביצועים גבוהים המשולבת באופן טבעי עם שירות אחסון פשוט של אמזון (S3). אנא עיין בהוראות מ הגדר ערוץ קלט נתונים לשימוש ב- Amazon FSx עבור Luster להדרכה על הגדרת מערכת קבצים FSx Luster כערוץ קלט נתונים.

3. התחל את עבודות ההדרכה

שלב זה מניח שכבר עשית זאת שינית את סקריפט האימון שלך והכין את מערך הנתונים כפי שהוזכר בסעיפים הקודמים. ל אפשר מקביליות נתונים מרוסקים, פשוט הגדר את sharded_data_parallel_degree ב PyTorch Estimator. במדריך זה, קבענו sharded_data_parallel_degree=128 ו instace_count=32 עבור צמתים p4d.24xlarge, מה שמצביע על כך שמצבי המודל יתפצלו על פני 128 GPUs מתוך סך 256 GPUs. בהתבסס על הערך הנבחר הזה, SMP יגדיר אוטומטית את מידת המקבילה של הנתונים ל-2 (מכיוון ש-256/128=2), כלומר יהיו לנו שני העתקים עבור מקביליות נתונים. כלל כללי לבחירת ערך אידיאלי עבור sharded_data_parallel_degree הוא להוסיף עוד צומת אחד לקבוצת השיתוף לכל 3B של פרמטרים של מודל. במדריך זה, גודל הדגם שלנו הוא 30B, אז עלינו להשתמש לפחות ב-10 צמתים לריסוק. ומכיוון ש-16 צמתים (128 GPUs) הם הכוח הקטן ביותר של 2 מעל הסף, קבענו sharded_data_parallel_degree=128.

עבור נקודות ביקורת, אנו מספקים גם קבוצה של כלי עזר למחסומים sharded_data_parallel_checkpoint.py , כולל כלי עזר לשחזור המלא state_dict למקרי שימוש מתקדמים. לבסוף, נוכל להשיק משרת הדרכה מבוזרת על ידי קריאה ל-fit() ב-Estimator.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. מעקב אחר עבודות ההדרכה

אתה יכול לגשת ליומני האימון ולעקוב אחר ניצול ה-GPU והזיכרון אמזון CloudWatch. הקפד להסתכל ביומנים של "algo-1" כי זה הצומת הראשי שלזרם הפלט שלו יש את יומני עבודות ההדרכה מכל המופעים.

השוואת ביצועים

סימנו מקביליות נתונים מרוסקים בספריית SMP בשני 16 ו-32 צמתים p4d.24xlarge עבור אורך רצף 512 ו-2048, בהתאמה. דגם GPT30 בעל פרמטר 2B מוגדר להשתמש ברוחב נסתר של 7168, 48 שכבות ו-64 ראשים. אתה יכול לאמץ את אותה תצורה בדיוק שבה אורך הרצף הוא 2048 על ידי הגדרה model_config = "gpt2-30b" במחברת ההדרכה. עם הגדרה זו, SMP השיג 73.52 דגימות לשנייה, מהירות של 39.7% בהשוואה ל-DeepSpeed ​​ZeRO-3. אם גודל האסימון שלך הוא 500 מיליארד, האצה זו פירושה כמעט 367 שעות של חיסכון בצמתים p4d.24xlarge, שווה ערך ליותר מ-$12,000 תקציב שנחסך בכל אימון! הטבלה הבאה מסכמת את תוצאות ההשוואה שלנו.

1/ עבור כל תצורת דגם, בדקנו תכונות, שלבים ותצורות שונות ב-DeepSpeed ​​ZeRO ובחרנו את זו שמספקת את התפוקה הטובה ביותר בתור קו הבסיס של DeepSpeed. המדד הופעל ענן מחשוב אלסטי של אמזון (אמזון EC2). 2/ תוצאות אלו מסתמכות על קולקטיבי תקשורת משופרים המותאמים ל-AWS אשר יהיו זמינים בקרוב. 3/ זמן הרכבה מוקרן ממהירות על סמך מספר האסימונים שעובדו.

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

סיכום

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

על המחברים

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

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

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

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

ארין הו הוא מנהל מוצר עבור AWS Deep Learning. היא עובדת על מוצרים שמקלים על הלקוחות להכשיר מודלים של למידה עמוקה ב-AWS. לעבודה מהנה בחוץ, היא נהנית מטיולים וסקי.