تدريب النماذج العملاقة بمقياس خطي تقريبًا باستخدام توازي البيانات المجزأة على Amazon SageMaker

في السعي لتحقيق الدقة الفائقة ، نمت نماذج التعلم العميق في مجالات مثل معالجة اللغة الطبيعية ورؤية الكمبيوتر بشكل كبير في الحجم في السنوات القليلة الماضية ، والتي يتم عدها في كثير من الأحيان من عشرات إلى مئات المليارات من المعلمات. يعد تدريب هذه النماذج العملاقة أمرًا صعبًا ويتطلب استراتيجيات توزيع معقدة. يبحث علماء البيانات ومهندسو التعلم الآلي باستمرار عن أفضل طريقة لتحسين حوسبة التدريب الخاصة بهم ، لكنهم يعانون من عبء الاتصال الذي يمكن أن يزداد جنبًا إلى جنب مع الحجم الكلي للكتلة.

هذا هو سبب إطلاقنا مؤخرًا sتوازي البيانات المتصلب on الأمازون SageMaker، وهي تقنية تدريب موزعة جديدة لحفظ الذاكرة في مكتبة نموذج SageMaker المتوازية (SMP). تم تصميم توازي البيانات المُقسمة خصيصًا للنماذج واسعة النطاق ويستخدم Amazon داخليًا ميكس التكنولوجيا تحت الغطاء ، جهد علمي لتقليل نطاق الاتصال عن طريق خفض تكلفة الاتصالات المكلفة المتجذرة في جمع المعلمات ومزامنة التدرج. مع نموذج GPT-30 2B مع طول التسلسل 2048 ، حققت هذه الميزة الجديدة 141 TFLOPs ، سرعة 39.7٪ مقارنة بـ DeepSpeed ​​ZeRO-3. بالنسبة إلى طراز 10B GPT-2 بطول التسلسل 512 ، حققت هذه الميزة الجديدة أيضًا 564 عينة في الثانية ، بزيادة قدرها 13.9٪ مقارنةً بـ PyTorch's Fully Sharded Data Parallel (FSDP). تذكر أنه في تدريب النموذج الضخم ، تُترجم كل نسبة تسريع إلى دولارات يتم توفيرها وزيادة الإنتاجية المكتسبة في فريقك.

في منشور المدونة هذا ، سنلقي أولاً نظرة فاحصة على عوامل التفاضل الرئيسية لتوازي البيانات المُقسمة ومتى نستخدمها. بعد ذلك ، ستتعلم كيفية تدريب نموذج GPT-30 بمعامل 2B على SageMaker بسهولة مع هذه الميزة الجديدة. أخيرًا ، سنقارن الأداء بخيارات أخرى مفتوحة المصدر ، لا سيما التفوق على DeepSpeed ​​ZeRO بنسبة تصل إلى 39.7٪ على 256 وحدة معالجة رسومات.

كيف يعمل توازي البيانات المجزأة ومتى يتم استخدامه

قبل أن نقدم توازي البيانات المجزأة ، دعنا نلقي نظرة على عائلة تقنياتها الأوسع. انتقلت مناهج التدريب الموزعة الحديثة للنماذج الكبيرة إلى نموذج حيث يتم مشاركة معلمات النموذج والتدرجات وحالات المحسن عبر العقد المتوازية للبيانات. على عكس Pipeline Parallelism الذي يتميز بالتعقيد الفطري في اختيار الطبقات للتقسيم عبر الأجهزة خاصةً عندما لا يدعم إطار العمل الخاص بك تقسيم النموذج الآلي، يحافظ هذا النموذج بأناقة على بساطة توازي البيانات ، مع إزالة قيود توازي البيانات حيث يجب أن يتناسب النموذج مع وحدة معالجة رسومات واحدة.

في الأطر الحالية التي تندرج تحت هذا النموذج ، ولا سيما DeepSpeed ​​ZeRO-3 و PyTorch's FSDP الذي تم تطويره من FairScale ، يتم تقسيم حالات النموذج عبر من جميع وحدات معالجة الرسومات (GPU) ، وهي إستراتيجية تقلل من استهلاك الذاكرة لكل وحدة معالجة رسومات على حساب تكبد نفقات اتصال كبيرة تزيد مع حجم المجموعة وبالتالي تتسبب في انخفاض قابلية التوسع بشكل كبير على نطاق واسع. في المقابل ، فإن توازي البيانات المجزأة في نموذج أقسام مكتبة SMP ينص على ملف مدركين للقياس بطريقة تقسيم كل نسخة طبق الأصل من حالات النموذج فقط داخل مجموعة فرعية من وحدات معالجة الرسومات.

دعونا نلقي نظرة فاحصة على نموذج التقسيم الواعي في MiCS ، التكنولوجيا الأساسية وراء البيانات المتوازية المتفرقة. يتمثل الحدس الكامن وراء هذا التصميم في أن تقسيم حالات التدريب عبر المجموعة الموازية للبيانات بأكملها قد لا يكون مطلوبًا لتدريب نموذج يحتوي على عشرات المليارات من المعلمات. على سبيل المثال ، 8 وحدات معالجة رسومات (GPU) V100 (32 جيجابايت لكل منها) كافية للاحتفاظ بنسخة متماثلة لحالات النموذج لنموذج معلمة 10B والذي يحتاج إلى حوالي 200 جيجابايت من الذاكرة عند التدريب باستخدام مُحسِّن Adam باستخدام دقة مختلطة. من خلال تحديد نسخة متماثلة كاملة من حالات النموذج في ملف أصغر مجموعة فرعية من وحدات معالجة الرسومات ، يمكننا بشكل فعال تقليل حجم الحمل الزائد مقارنة بـ DeepSpeed ​​و PyTorch FSDP. تعمل البيانات المتوازية على الاستفادة من التقنيات الأخرى في نظام MiCS مثل الاتصال الهرمي ومزامنة التدرج ثنائي القفزات. لمزيد من المعلومات ، تحقق من التحجيم شبه الخطي لتدريب النموذج العملاق على AWS or MiCS: التحجيم شبه الخطي لتدريب النموذج العملاق على السحابة العامة.

الآن ، كيف تعرف متى تختار البيانات المجزأة الموازية على تقنيات التدريب الموزعة الأخرى؟ القاعدة العامة هي أنه إذا كان نموذجك يحتوي على أقل من مليار معلمة ويمكن أن يتناسب مع ذاكرة وحدة معالجة الرسومات ، مكتبة SageMaker البيانات المتوازية or مترجم تدريب SageMaker يمكن أن يكون كافيا بالنسبة لك. إذا كان لديك لغة أو نماذج رؤية كمبيوتر أكبر ، فإن اقتراحنا هو تدريبها باستخدام تقنية موازاة البيانات المُقسمة جنبًا إلى جنب مع تفعيل نقاط التفتيش و تفعيل التفريغ في المكتبة الموازية لنموذج SageMaker أولاً ، قبل التقنيات الأخرى مثل التوازي الموتر أو موازاة خطوط الأنابيب.

استخدام توازي البيانات المُقسمة لتدريب GPT-2 على Amazon SageMaker

دعنا الآن نتعلم كيفية تدريب نموذج GPT-2 مع بيانات مجزأة متوازية ، مع SMP الذي يغلف التعقيد من أجلك. هذه دفتر تعليمي كامل يرشدك خلال العملية برمتها ، من معالجة البيانات ، وتحديد وتقديم وظائف التدريب ، إلى مراقبة سجلات التدريب. فيما يلي نظرة عامة موجزة تسلط الضوء على الخطوات الرئيسية لاستخدام هذه الميزة.

1. ابدأ

توازي البيانات المشتركة متاح في PyTorch v1.12.0 + ويعمل مع كل من FP16 و BF16. أسهل طريقة لاستخدام مكتبة SMP هي من خلال حاوية AWS Deep Learning Container سابقة الإنشاء لـ PyTorch. ومع ذلك ، إذا كنت تريد إحضار حاوية Docker الخاصة بك ، فيمكنك الرجوع إلى قم بإنشاء حاوية Docker الخاصة بك باستخدام مكتبة SageMaker الموزعة للنموذج المتوازي. للبدء ، اتبع تعديل البرنامج النصي للتدريب على PyTorch لتكييف واجهات برمجة تطبيقات SMPs في نص التدريب الخاص بك. في هذا القسم ، نستدعي فقط بضع خطوات رئيسية مع مقتطفات التعليمات البرمجية من البرنامج النصي للتدريب الجاهز للاستخدام train_gpt_simple.py. يمكنك متابعة التعليقات في البرنامج النصي و وثيقة 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 من أجل اللمعان للحصول على إرشادات حول إعداد نظام ملفات FSx Luster كقناة إدخال بيانات.

3. ابدأ وظائف التدريب

تفترض هذه الخطوة أن لديك بالفعل تعديل نص التدريب الخاص بك وأعد مجموعة البيانات كما هو مذكور في الأقسام السابقة. إلى تمكين توازي البيانات المجزأة ، ما عليك سوى تعيين sharded_data_parallel_degree في ال مقدر PyTorch. في هذا البرنامج التعليمي ، حددنا sharded_data_parallel_degree=128 و  instace_count=32 بالنسبة للعقد p4d.24xlarge ، مما يشير إلى أن حالات النموذج ستتم مشاركتها عبر 128 وحدة معالجة رسومات من بين إجمالي 256 وحدة معالجة رسومات. بناءً على هذه القيمة المحددة ، سيقوم SMP تلقائيًا بتعيين درجة موازية البيانات إلى 2 (لأن 256/128 = 2) ، مما يعني أنه سيكون لدينا نسختان متماثلتان لتوازي البيانات. قاعدة عامة لاختيار القيمة المثالية لـ sharded_data_parallel_degree هو إضافة عقدة أخرى إلى مجموعة المشاركة لكل 3B من معلمات النموذج. في هذا البرنامج التعليمي ، حجم نموذجنا هو 30B ، لذلك يجب أن نستخدم 10 عقد على الأقل للتجزئة. ونظرًا لأن 16 عقدة (128 وحدة معالجة رسومات) هي أصغر قوة من 2 أعلى من العتبة ، قمنا بتعيين sharded_data_parallel_degree=128.

بالنسبة إلى نقاط التفتيش ، نوفر أيضًا مجموعة من أدوات نقاط التفتيش في sharded_data_parallel_checkpoint.py ، بما في ذلك أداة لإعادة بناء كامل state_dict لحالات الاستخدام المتقدمة. أخيرًا ، يمكننا إطلاق وظيفة تدريبية موزعة عن طريق استدعاء fit () على المقدر.

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 من احتواء حجم دفعة أكبر وتقليل مستوى تراكم التدرج المطلوب لملاءمة حجم دفعة عالمية معين.

وفي الختام

في هذا المنشور ، قدمنا ​​أسلوبًا جديدًا للتدريب الموزع - توازي البيانات المُقسمة - وكيف يُسرع من تدريب النموذج العملاق مع التحجيم الخطي القريب على Amazon SageMaker. استعرضنا أيضًا كيفية تدريب نموذج GPT-2 باستخدام التقنية الجديدة التي تتبع ذلك مثال كامل. يمكنك متابعة أمثلة على Amazon SageMaker GitHub repo لتتبع جميع الأمثلة الموازية لنماذج SageMaker أو حضورنا التالي وزعت ورش عمل تدريبية. لمعرفة المزيد حول توازي البيانات المجزأة ، يرجى الاطلاع على توثيق.

عن المؤلفين

إميلي ويبر انضم إلى AWS بعد إطلاق SageMaker مباشرة ، وما زال يحاول إخبار العالم عنها منذ ذلك الحين! خارج بناء تجارب تعلم جديدة للعملاء ، تتمتع إميلي بالتأمل ودراسة البوذية التبتية.

يمكن Karakus هو عالم تطبيقي أقدم في AWS ، يعمل على تحسين التعلم العميق الموزع على نطاق واسع على AWS. تغطي اهتماماته البحثية التعلم العميق ، والتحسين الموزع ، والأنظمة الموزعة ، ونظرية المعلومات. خارج العمل ، يستمتع بركوب الدراجات والسفر والقراءة والتعلم.

راهول هويلغول هو مهندس برمجيات أول في AWS. إنه يعمل على أنظمة التعلم العميق الموزعة ، من أجل جعل تدريب نماذج التعلم العميق الكبيرة في السحابة أمرًا سهلاً وفعالاً. في أوقات فراغه ، يستمتع بالتصوير وركوب الدراجات والبستنة.

سوهيت كودجولي مهندس تطوير برمجيات مع مجموعة AWS للذكاء الاصطناعي تعمل على أطر التعلم العميق. في أوقات فراغه ، يستمتع بالتنزه والسفر والطهي.

ايرين هو هو مدير منتج لـ AWS Deep Learning. تعمل على منتجات تسهل على العملاء تدريب نماذج التعلم العميق على AWS. من أجل المتعة خارج العمل ، تستمتع برياضة المشي لمسافات طويلة والتزلج.