Scaling & Efficiency2019متقدم12 دقيقة قراءة

Megatron-LM: تدريب نماذج لغوية بمليارات المعاملات باستخدام التوازي على مستوى النموذج

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Shoeybi, M. · Patwary, M. · Puri, R. · LeGresley, P. · Casper, J. · Catanzaro, B. — arXiv

المشكلة

بحلول 2019 كانت المعادلة واضحة: كلما كبر نموذج تحسّنت نتائجه في مهام اللغة. GPT-2 وصل إلى 1.5 مليار معامل، والمؤشرات كلها كانت تقول إن مزيداً من التوسّع سيُعطي أداءً أفضل. المشكلة أن المعالج الرسومي الواحد لا يملك ذاكرة كافية — 16 أو 32 غيغابايت في العادة — بينما نموذج بمليارات المعاملات يحتاج إلى تخزين الأوزان والتدرّجات وحالات المُحسِّن والتنشيطات، وهذا يتجاوز 40 غيغابايت بسهولة. لا يحل شيئاً هنا لأنه ينسخ النموذج كاملاً على كل معالج — فإذا لم يتسع النموذج في معالج واحد فلن يتسع في أيٍّ منها. والحلول الموجودة آنذاك مثل GPipe كانت تحتاج مُصرِّفات مخصصة أو إعادة بناء لأُطر العمل البرمجية، فلم يكن تبنّيها سهلاً.

الإسهام

طريقة فعّالة لتوازي الموتّرات داخل الطبقة الواحدة، مُصمَّمة لبنية المحوِّل تحديداً ولا تحتاج أي مُصرِّف مخصص — كل ما تفعله هو إدخال بضع عمليات اتصال في شيفرة PyTorch العادية. الفكرة الجوهرية هي تقسيم مصفوفات الأوزان في كل طبقة — سواء في الشبكة الأمامية أو الانتباه متعدد الرؤوس — وفق نمط «أعمدة ثم صفوف»، فيحتفظ كل معالج رسومي بشريحة من كل طبقة ولا يحتاج سوى عملية تجميع واحدة لكل كتلة فرعية. بهذا الأسلوب مع بدقة مختلطة، دُرِّب نموذج GPT-2 بحجم 8.3 مليار معامل على 512 معالجاً رسومياً وبلغت الإنتاجية 15.1 بيتا فلوبس بكفاءة توسّع 76%. كذلك أثبت الباحثون أن نقل تسوية الطبقة إلى موضع مختلف في بنية BERT شرط أساسي للتوسّع إلى ما بعد BERT-Large، وحقق نموذجهم BERT ذو 3.9 مليار معامل أفضل نتيجة على معيار RACE.

الأثر

تحوّلت Megatron-LM إلى البنية التحتية المعيارية لتدريب النماذج اللغوية الكبيرة داخل NVIDIA وخارجها. تقريباً كل منظومة تدريب واسعة النطاق ظهرت بعدها تبنّت أسلوبها في توازي الموتّرات — من تدريب GPT-3 إلى BLOOM وغيرها من نماذج المليارات. ولمّا اجتمعت مع مُحسِّن ZeRO من DeepSpeed نشأ ما يُعرف بالتوازي ثلاثي الأبعاد (بيانات + موتّرات + خط أنابيب)، وهو الأسلوب الذي تقوم عليه أضخم النماذج اليوم. ولا تزال قاعدة الشيفرة المفتوحة المصدر واحدةً من أكثر الأُسس انتشاراً في للنماذج اللغوية الكبيرة.

تخيّل مطبخ مطعم يُجهّز وليمة كبيرة. توازي البيانات يُشبه أن يكون لديك عشرة مطابخ متطابقة، كلٌّ منها يطهو القائمة كاملة من مكوّناته — سريع، لكن كل مطبخ يحتاج نسخة كاملة من كل المعدات. يُشبه خط إنتاج: محطة تُقطّع، والتالية تقلي، والثالثة تُرتّب الأطباق — سريع أيضاً، لكن كل محطة تنتظر التي قبلها حتى تنتهي.

في Megatron-LM يختلف تماماً: بدلاً من تقسيم المحطات، يُقسّم العمل داخل المحطة الواحدة. تخيّل أن أربعة طهاة يقفون عند طاولة التقطيع نفسها، كلٌّ منهم يُقطّع ربع الخضروات في الوقت ذاته. كل ما يحتاجونه هو أن يجمعوا أكوامهم مرة واحدة قبل تمريرها للمحطة التالية. أقل انتظار ممكن، وأقصى عمل متوازٍ — ودون أي إعادة تصميم للمطبخ.

جدار الذاكرة: لماذا لا يكفي توازي البيانات

في توازي البيانات، يُنسخ النموذج بالكامل على كل . كل وحدة تأخذ دفعة جزئية مختلفة وتحسب التدرّجات، ثم تتبادل الوحدات جميعها هذه التدرّجات عبر عملية . الأسلوب ممتاز — لكنه يتوقف عن العمل لحظة أن يصبح النموذج أكبر من أن يتسع في وحدة واحدة.

لنضع أرقاماً ملموسة: نموذج بملياري معامل بدقة FP16 يستهلك 4 غيغابايت للأوزان وحدها. أضف حالات المُحسِّن بدقة FP32 — يحتفظ بعزمَين لكل معامل — ثم التدرّجات والتنشيطات، وستجد أن المعالج الواحد يحتاج أكثر من 40 غيغابايت. هذا يتجاوز ذاكرة معظم المعالجات الرسومية التي كانت متاحة عام 2019.

توازي خط الأنابيب — كما في GPipe — يُقسّم النموذج إلى مراحل متتابعة على معالجات مختلفة. لكنه يخلق ما يُعرف بـ«فقاعة خط الأنابيب»: المعالجات تبقى عاطلة بانتظار التنشيطات من المراحل السابقة، وكلما ازداد عمق الخط ازدادت الفقاعة سوءاً.

المطلوب إذاً كان شيئاً مختلفاً: طريقة لتقسيم كل طبقة عبر المعالجات، بحيث تشارك كل وحدة في حساب كل طبقة في الوقت نفسه. هذا بالضبط ما يفعله توازي الموتّرات — وMegatron-LM هو من حوّله إلى واقع عملي لبنية المُحوِّل.

افتح في المختبر

قارن بين الاستراتيجيات الثلاث: توازي البيانات ينسخ النموذج بالكامل، توازي خط الأنابيب يقسّمه طبقةً طبقة، وتوازي الموتّرات يقسّم الطبقة نفسها من الداخل.

تستيقظ التجربة عند وصولك…

الفكرة الجوهرية: تقطيع الشبكة الأمامية عبر الوحدات

الشبكة الأمامية في المُحوِّل تتكوّن من طبقتين خطيتين تفصل بينهما GeLU. العملية الحسابية هي $Y = \text{GeLU}(XA)$ ثم $Z = \text{Dropout}(YB)$ . المصفوفتان $A$ و $B$ هما أكبر ما يستهلك الذاكرة. السؤال الآن: كيف نُقسّم هاتين المصفوفتين على عدة معالجات؟

الخيار الأول: تقسيم $A$ أفقياً (على طول الصفوف). كل معالج يأخذ شريحة أفقية من $A$ مع الجزء المقابل من $X$ . المشكلة هنا أن GeLU دالة لاخطية، أي أن: $\text{GeLU}(X_1 A_1 + X_2 A_2) \neq \text{GeLU}(X_1 A_1) + \text{GeLU}(X_2 A_2)$ . هذا يفرض إجراء تجميع شامل قبل تطبيق GeLU — أي خطوة اتصال إضافية في منتصف الحساب، وهذا مكلف.

الخيار الثاني (ما اختاره Megatron): تقسيم $A$ عمودياً (على طول الأعمدة). كل معالج يأخذ شريحة عمودية: $A = [A_1, A_2]$ . الفائدة أن كل معالج يستطيع أن يحسب $\text{GeLU}(X A_i)$ باستقلال تام، لأن GeLU تُطبَّق على كل جزء وحده دون الحاجة لنتائج المعالجات الأخرى. لا مزامنة مطلوبة حتى النهاية. بعدها تُقسَّم $B$ على طول صفوفها فتستقبل المخرجات المُقسَّمة مباشرة. المزامنة الوحيدة هي عملية تجميع شامل واحدة بعد الضرب المصفوفي الثاني، مباشرة قبل .

[Y_1, Y_2] = [\text{GeLU}(XA_1), \;\text{GeLU}(XA_2)]

الشبكة الأمامية بتوازي الأعمدة — تطبيق GeLU بشكل مستقل لكل جزء — عندما نُقسّم

A

عمودياً، يحسب كل معالج

\text{GeLU}(XA_i)

باستقلال تام. هذا يُلغي نقطة المزامنة قبل الدالة اللاخطية — وهي الملاحظة المحورية التي تجعل أسلوب Megatron فعّالاً.

افتح في المختبر

شاهد كيف تُوزَّع مصفوفات الأوزان في الشبكة الأمامية على المعالجات. التقسيم العمودي لـ A يتجنّب المزامنة قبل GeLU، والتقسيم الصفي لـ B لا يحتاج سوى تجميع شامل واحد في النهاية.

تستيقظ التجربة عند وصولك…

نمط الاتصال بسيط وأنيق. كل ما في الأمر مُعاملان مخصصان: $f$ يتصرف كدالة هُوية في وينفّذ تجميعاً شاملاً في . $g$ يعمل بالعكس تماماً: تجميع شامل في التمرير الأمامي وهُوية في التمرير العكسي. هذان المُعاملان متعاكسان — معاً يضمنان صحة المخرجات في الأمام وصحة التدرّجات في الخلف، بعملية تجميع شامل واحدة فقط لكل اتجاه في كل كتلة. والتنفيذ بأكمله لا يتجاوز بضعة أسطر من PyTorch.

الانتباه: تصميم طبيعي للتوازي

متوازٍ بطبيعته: كل رأس يعمل على إسقاط مُتعلَّم مختلف للمدخلات، أي أن الرؤوس حسابات مستقلة لا تحتاج أن تُدمج إلا في الخطوة الأخيرة. Megatron تستغل هذه الخاصية بتقسيم مصفوفات الإسقاط Q وK وV على طول أعمدتها، فيحسب كل معالج مجموعة فرعية من الرؤوس.

مثال ملموس: نموذج فيه 16 رأس انتباه و4 معالجات رسومية — كل معالج يتولى 4 رؤوس. مصفوفات Q وK وV مقسَّمة بحيث يحصل كل معالج على الأجزاء الخاصة برؤوسه فقط. بعد حساب ، يكون لدى كل معالج مخرج جزئي. مصفوفة الإسقاط الخطي $O$ تُقسَّم صفياً — تماماً مثل $B$ في الشبكة الأمامية — فتدخل المخرجات الجزئية مباشرة دون مزامنة، ولا يلزم سوى تجميع شامل واحد بعد إسقاط المخرج.

الأنيق في هذا التصميم أن كتلة الانتباه وكتلة الشبكة الأمامية تستخدمان نمط الاتصال نفسه: مُعامل $f$ عند المدخل ومُعامل $g$ عند المخرج. المحصلة لطبقة محوِّل كاملة هي عمليتا تجميع شامل في التمرير الأمامي وعمليتان في التمرير العكسي — لا أكثر.

افتح في المختبر

شاهد كيف تُوزَّع رؤوس الانتباه على المعالجات. كل معالج يحسب مجموعته من الرؤوس باستقلال تام، ولا يلزم سوى تجميع شامل واحد بعد إسقاط المخرج.

تستيقظ التجربة عند وصولك…

المُعاملان f وg: سطران من الشيفرة يُمكّنان كل شيء

أناقة Megatron تكمن في مُعاملَي تفاضل تلقائي فقط. فكّر فيهما كبوابتَين عند حدود كل كتلة متوازية:

المُعامل $f$ : في التمرير الأمامي يُمرّر المدخل كما هو (دالة هُوية). في التمرير العكسي يُجري تجميعاً شاملاً ليجمع التدرّجات من كل المعالجات. يُوضع قبل الكتلة المتوازية — كل معالج يتلقى المدخل كاملاً ويحسب حصته منه.

المُعامل $g$ : في التمرير الأمامي يُجري تجميعاً شاملاً لدمج المخرجات الجزئية من كل المعالجات. في التمرير العكسي يُمرّر التدرّجات كما هي. يُوضع بعد الكتلة المتوازية ليجمع النتائج الجزئية في مخرج واحد كامل.

$f$ و $g$ زوج متعاكس، وهما كل ما يلزم من اتصال. لا مُصرِّف مخصص ولا إطار عمل خاص — فقط فئتا autograd.Function في PyTorch.

المُعاملان f وg — طبقة الاتصال الكاملة لـ Megatron في PyTorchpython

مبسَّط لإظهار الفكرة — ليس التنفيذ الحقيقي.

import torch
import torch.distributed as dist

class f(torch.autograd.Function):
    """أمامي: دالة هُوية. عكسي: تجميع شامل للتدرّجات."""
    @staticmethod
    def forward(ctx, x):
        return x                            # مرور بلا تغيير

    @staticmethod
    def backward(ctx, grad):
        dist.all_reduce(grad)                # اجمع التدرّجات من كل الوحدات
        return grad

class g(torch.autograd.Function):
    """أمامي: تجميع شامل للمخرجات الجزئية. عكسي: دالة هُوية."""
    @staticmethod
    def forward(ctx, x):
        dist.all_reduce(x)                   # اجمع النتائج الجزئية من كل الوحدات
        return x

    @staticmethod
    def backward(ctx, grad):
        return grad                          # مرور بلا تغيير

# الاستخدام في شبكة أمامية بتوازي الموتّرات:
# x = f.apply(x)           # كل وحدة تتلقى المدخل الكامل
# y = gelu(x @ A_i)        # كل وحدة تحسب جزءها العمودي
# z = y @ B_i              # كل وحدة تحسب جزءها الصفي
# z = g.apply(z)           # تجميع شامل ← المخرج الكامل

نتائج التوسّع: 512 وحدة معالجة، كفاءة 76%

درّب الباحثون نماذج GPT-2 بأحجام متصاعدة: 1.2 و2.5 و4.2 و8.3 مليار معامل. في كل الحالات استخدموا توازي موتّرات ثُماني الاتجاه داخل كل عقدة DGX — أي 8 معالجات متصلة بوصلة NVLink — مع توازي بيانات بين العقد.

النموذج الأكبر (8.3 مليار معامل) على 512 معالجاً حقق مستدامة بلغت 15.1 بيتا ، أي 76% مقارنة بخط أساس قوي على معالج واحد بلغ 39 تيرا فلوبس (30% من الأداء الذروي لـ V100). في سياق 2019 كان هذا إنجازاً استثنائياً: أغلب الأساليب الأخرى كانت تفقد كثيراً من كفاءتها عند هذا الحجم.

القرار الهندسي المفصلي كان حصر مجموعات توازي الموتّرات داخل نطاق NVLink الواحد (أي داخل خادم واحد). وصلة NVLink توفر عرض نطاق ترددي يصل إلى 300 غيغابايت/ثانية بين المعالجات في الخادم الواحد، بينما الوصلات بين العقد أبطأ بـ 10–25 مرة. الفكرة بسيطة: التوازي الذي يحتاج اتصالاً مكثفاً (توازي الموتّرات) يبقى في الوصلة السريعة، والتوازي الأخف اتصالاً (توازي البيانات) يمتد بين العقد. بهذا حققت Megatron توسّعاً شبه خطي.

افتح في المختبر

استكشف كيف تنمو الإنتاجية مع زيادة عدد المعالجات. توازي موتّرات ثُماني داخل العقدة مع توازي بيانات بين العقد يحقق كفاءة 76% على 512 معالجاً.

تستيقظ التجربة عند وصولك…

إصلاح BERT: إعادة ترتيب تسوية الطبقة للتوسّع

توسيع BERT إلى ما بعد 336 مليون معامل كان مشكلة معروفة: التدريب كان ينهار والأداء يتراجع بدلاً من أن يتحسّن. باحثو Megatron-LM شخّصوا السبب الجذري: في بنية BERT الأصلية تُطبَّق بعد (وضع post-norm). المشكلة أن القيم في المسار المتبقي تتراكم وتكبر مع كل طبقة، وفي النماذج العميقة ينتهي الأمر بعدم استقرار التدريب.

الحل كان بسيطاً من حيث التنفيذ لكنه حاسم: نقل تسوية الطبقة إلى قبل كتلتي الانتباه الذاتي والشبكة الأمامية (وضع pre-norm)، كما يفعل GPT-2. بهذا الترتيب يبقى المسار المتبقي نظيفاً — يحمل قيماً خاماً تُجمع خطياً — والتسوية تُطبَّق على مدخلات كل طبقة فرعية فتبقى التنشيطات مضبوطة.

بهذا التعديل نجح الباحثون في تدريب نماذج BERT تصل إلى 3.9 مليار معامل — أي 12 ضعف حجم BERT-Large — وحققوا أفضل النتائج على معيار RACE لفهم القراءة بدقة 90.9%.

افتح في المختبر

قارن بين وضع post-norm في BERT الأصلي ووضع pre-norm في إصلاح Megatron-LM. لاحظ الفرق في سلوك التنشيطات عبر 24 طبقة.

تستيقظ التجربة عند وصولك…

نتائج المعايير: نماذج أكبر بأداء أفضل

النتائج أظهرت نمطاً واضحاً: كلما كبر النموذج تقارب أسرع ووصل إلى حيرة أقل. نموذج 8.3 مليار معامل حقق أفضل النتائج على عدة معايير مرجعية:

WikiText-103: حيرة 10.8 مقابل 15.8 سابقاً — تحسّن نسبي بنسبة 32%.
LAMBADA: دقة 66.5% مقابل 63.2% — ما يعكس تحسّناً في الفهم اللغوي بعيد المدى.
RACE (فهم القراءة): دقة 90.9% بنموذج BERT ذي 3.9 مليار معامل، متفوقاً على الرقم السابق 89.4%.

هذه النتائج أكدت فرضية محورية: متى توفّرت قدرة حسابية كافية واستراتيجية توازٍ فعّالة، فإن مجرد تكبير نماذج المُحوِّل يُحقق قفزات حقيقية في الجودة — وهذا بالضبط المبدأ الذي رسم ملامح الحقبة التالية مع GPT-3 وما بعده.

الإرث: كتاب قواعد التدريب الموزَّع

2018
GPipe (Google) — توازي خط الأنابيب
يقسّم النموذج إلى مراحل متتابعة على معالجات مختلفة. فعّال لكنه يعاني من فقاعات خط الأنابيب ويحتاج إطار عمل مخصصاً.
2019
Megatron-LM — توازي الموتّرات
توازي داخل الطبقة يقسّم مصفوفات الأوزان في كل طبقة محوِّل. يعمل بـ PyTorch الأصلي مباشرة، وبلغ 8.3 مليار معامل بكفاءة توسّع 76% على 512 معالجاً.
2019
ZeRO (DeepSpeed) — تقسيم حالات المُحسِّن
يُوزّع حالات المُحسِّن والتدرّجات والمعاملات على رتب توازي البيانات. مستقل عن توازي الموتّرات، لكنهما يتكاملان بقوة.
2020
GPT-3 (175 مليار معامل)
دُرِّب بتوازي موتّرات على طراز Megatron مقترناً بتوازي بيانات. أثبت أن الوصول إلى 175 مليار معامل يفتح الباب أمام التعلّم ببضعة أمثلة.
2021
عصر التوازي ثلاثي الأبعاد
جمع Megatron-DeepSpeed بين توازي الموتّرات وخط الأنابيب والبيانات في إطار واحد. هذا «التوازي ثلاثي الأبعاد» صار الوصفة المعتمدة لتدريب نماذج بمئات المليارات فما فوق.
2022
BLOOM (BigScience، 176 مليار)
أكبر نموذج لغوي مفتوح المصدر في حينه، دُرِّب بأدوات Megatron-DeepSpeed على 384 معالج A100. وريث مباشر لتقنيات Megatron-LM.

أثر Megatron-LM يتخطى أي نموذج بعينه — فقد أسست ما يمكن تسميته «كتاب وصفات التدريب الموزَّع» للنماذج اللغوية الكبيرة: توازي الموتّرات داخل العقدة حيث الوصلة السريعة، توازي خط الأنابيب عبر عدد محدود من العقد، وتوازي البيانات عبر العنقود بأكمله. كل منظومة تدريب رئيسية اليوم — من بنية Meta لتدريب LLaMA إلى إعدادات Google لتدريب PaLM — تسير على هذا النهج. والفكرة الأساسية — أن بضع عمليات تجميع شامل في PyTorch العادي تُغني عن مُصرِّفات مخصصة — هي ما فتح الباب لتدريب نماذج المليارات أمام الجميع.

المرجعShoeybi, Patwary, Puri, LeGresley, Casper, Catanzaro. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv, 2019.

مصطلحات هذه الورقة

توازي مصفوفات الموتّراتTensor Parallelism
توازي النموذجModel Parallelism
توازي البياناتData Parallelism
التوازي المتسلسل للطبقاتPipeline Parallelism
الاختزال الشاملAllreduce
كفاءة التوسّعScaling Efficiency
التدريب بالدقة المختلطةMixed Precision Training
معدل التدفق والإنتاجيةThroughput
العمليات الحسابية العائمةFLOPs
وحدة معالجة الرسومياتGPU
تجميع التدرجات الحسابيةGradient Accumulation
التسوية الطبقيةLayer Normalization
الوصلة التجاوزيةResidual Connection
وحدة الخطأ الخطية الغاوسية (GELU)GELU
المحوِّلTransformer