السَّمْت

Distributed Training2019متقدم12 دقيقة قراءة

ZeRO: تحسينات الذاكرة نحو تدريب نماذج بتريليون معامل

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Rajbhandari, S. · Rasley, J. · Ruwase, O. · He, Y. — SC

المشكلة

بحلول عام 2019، أثبتت النماذج التي تضمّ مليارات المعاملات قدرتها على تحقيق قفزات كبيرة في الدقة، لكن تدريبها اصطدم بسقف ذاكرة GPU. المشكلة أن التوازي في البيانات يُنشئ نسخة كاملة من النموذج على كل GPU — والتدرّجات والمعاملات — وبالتالي تتكرّر نفس البيانات على كل جهاز دون داعٍ. نموذج بحجم 1.5 مليار معامل مع مُحسِّن Adam يستهلك 24 جيجابايت لحالات النموذج وحدها، وهذا قبل حتى أن نحسب التنشيطات. البديل المتاح — التوازي في النموذج — يُوزّع الأجزاء على عدة أجهزة، لكنه يتطلب إعادة كتابة الشيفرة لكل بنية على حدة، ويعاني من كفاءة حوسبة منخفضة بسبب الاعتماديات التسلسلية بين الطبقات، ولا يتوسّع جيداً. ببساطة، لم يكن هناك حلّ يجمع بين وفرة الذاكرة التي يوفّرها التوازي في النموذج وسهولة الاستخدام وكفاءة الحوسبة التي يتميّز بها التوازي في البيانات.

الإسهام

يُزيل ZeRO (مُحسِّن عديم التكرار) البيانات المُكرَّرة من ذاكرة المتوازي عبر ثلاث مراحل تصاعدية: الأولى تُوزّع حالات المُحسِّن (توفير 4×)، والثانية تُوزّع التدرّجات أيضاً (توفير 8×)، والثالثة تُوزّع المعاملات ذاتها (توفير خطّي يتناسب مع عدد الأجهزة). إلى جانب ذلك، يُعالج ZeRO-R الذاكرة المتبقّية بتقسيم التنشيطات وتفريغها إلى CPU، واستخدام مخازن مؤقتة ثابتة الحجم، ونظام لإلغاء . المحصّلة: تدريب نماذج تصل إلى 100 مليار معامل على 400 GPU يبلغ 15 بيتافلوبس — أي زيادة 8 أضعاف في حجم النموذج القابل للتدريب وتحسُّن 10 أضعاف في الأداء مقارنة بأفضل ما سبقه.

الأثر

تحوّل ZeRO إلى العمود الفقري لمنصة DeepSpeed، وأعاد تشكيل طريقة تدريب النماذج اللغوية الكبيرة على مستوى الصناعة. تبنّت PyTorch أفكاره الجوهرية في FSDP، وصارت هذه الاستراتيجية هي الأسلوب المعتمد لتدريب نماذج بعشرات ومئات المليارات من المعاملات. وقت نشر الورقة، مكّن ZeRO من تدريب Turing-NLG بـ 17 مليار معامل، ومهّد الطريق مباشرة لمسار التوسُّع الذي أنتج GPT-3 وMegatron-Turing وما تلاهما. ثم توسّعت العائلة لتشمل ZeRO-Offload وZeRO-Infinity وZeRO++، فامتدّ نطاق العمل من عناقيد GPU الضخمة وصولاً إلى جهاز GPU واحد.

تخيّل فريقاً من الطهاة يُحضّرون الوصفة نفسها في مطبخ كبير. كل طاهٍ يحتفظ بنسخته الكاملة من كتاب الوصفات، ورفّ بهارات كامل، وطقم أكواب قياس كامل — مع أن الجميع يطبخ الطبق ذاته. النتيجة أن المطبخ يمتلئ بالنُّسخ وتتقلّص المساحة المتاحة للطبخ الفعلي.

هنا يأتي ZeRO بدور رئيس الطهاة ويقول: «مزّقوا كتاب الوصفات — كل طاهٍ يأخذ صفحاته فقط. وزّعوا رفّ البهارات — كل واحد يحتفظ ببهاراته فحسب. وإذا احتجتَ بهاراً ليس عندك، نادِ على من يملكه وسيُناوله لك.» بهذا الأسلوب، المطبخ نفسه يستطيع التعامل مع وصفة أعقد بعشر مرّات، لأن كل المساحة التي كانت مشغولة بالنُّسخ المُكرَّرة تحرّرت للطبخ الفعلي.

جدار الذاكرة: لماذا تنفد ذاكرة GPU

قبل أن نفهم ZeRO، لا بدّ أن نعرف ما الذي يملأ ذاكرة فعلياً أثناء التدريب. حين ندرّب نموذجاً بعدد معاملات Ψ باستخدام مع مُحسِّن ، تنقسم الذاكرة إلى فئتين رئيسيتين: حالات النموذج والحالات المتبقّية.

حالات النموذج تتكوّن من ثلاثة أجزاء. المعاملات بصيغة fp16 وتشغل 2Ψ بايت. بصيغة fp16 وتشغل 2Ψ بايت أخرى. ثم حالات المُحسِّن: يحتفظ Adam بنسخة fp32 من المعاملات (4Ψ بايت) إضافة إلى الزخم (4Ψ بايت) والتباين (4Ψ بايت) — أي 12Ψ بايت للمُحسِّن وحده. المجموع الكلّي لحالات النموذج يصل إلى 16Ψ بايت. بالنسبة لنموذج مثل GPT-2 بحجم 1.5 مليار معامل، هذا يعني 24 جيجابايت — ولم نحسب التنشيطات بعد.

أما الحالات المتبقّية فتشمل التنشيطات (قد تستهلك 60 جيجابايت لـ GPT-2 عند 32)، والمخازن المؤقتة (نحو 6 جيجابايت في النماذج الكبيرة)، الذاكرة التي تُنتج فراغات غير قابلة للاستخدام. كل هذه المكوّنات تتنافس مع حالات النموذج على ذاكرة GPU نفسها.

افتح في المختبر
استكشف توزيع ذاكرة GPU أثناء التدريب بالدقة المختلطة مع مُحسِّن Adam.
تستيقظ التجربة عند وصولك…

معضلة التوازي: الذاكرة مقابل الكفاءة

قبل ZeRO، كان المهندسون أمام مفاضلة صعبة بين أسلوبين للتوازي، كل منهما يتفوّق في جانب ويخسر في آخر.

ينسخ النموذج كاملاً على كل GPU. كل جهاز يعالج دُفعة مختلفة من البيانات، يحسب التدرّجات، ثم يُزامنها مع البقية عبر . الميزة أنه بسيط وفعّال — كل GPU يعمل بالتوازي — لكن الثمن هو هدر الذاكرة: كل جهاز يحمل نسخة كاملة من حالات النموذج (16Ψ بايت)، رغم أنها نسخ متطابقة تماماً.

يُوزّع أجزاء النموذج على عدة أجهزة، فكل GPU يحمل جزءاً فقط من المعاملات. هذا يحلّ مشكلة الذاكرة، لكنه يفتح مشكلات أخرى: تقسيم كل نموذج يتطلب إعادة كتابة الشيفرة يدوياً، والاعتماديات بين الطبقات تُنشئ اختناقات تسلسلية تترك بعض الأجهزة خاملة، ونقل التنشيطات بين المراحل يُضيف . في الممارسة، نادراً ما يتجاوز هذا الأسلوب 50% من كفاءة الحوسبة.

السؤال الذي طرحه ZeRO: ماذا لو قسّمنا حالات النموذج كما يفعل التوازي في النموذج، لكن أبقينا على نمط التنفيذ البسيط الخاص بالتوازي في البيانات؟

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

ZeRO-DP: ثلاث مراحل لتحرير الذاكرة

فكرة ZeRO الجوهرية بسيطة: في التوازي المعياري في البيانات، كل GPU يحتفظ بنسخة متطابقة من جميع حالات النموذج — وهذا تكرار محض لا مبرّر له. بدل ذلك، يُوزّع ZeRO هذه الحالات بين الأجهزة، ويستخدم عمليات التواصل الجماعي لتجميع الأجزاء حين الحاجة فقط.

الملاحظة الأساسية هنا أن كل جهاز لا يحتاج كل البيانات في كل لحظة. حالات المُحسِّن لا تُستخدم إلا عند تحديث المعاملات. التدرّجات لا تُستخدم إلا بعد . والمعاملات نفسها — وإن كانت مطلوبة أثناء التمريرين الأمامي والخلفي — يمكن جمعها مؤقتاً من الأجهزة الأخرى عند الحاجة.

يُطبّق ZeRO-DP هذا المبدأ عبر ثلاث مراحل تراكمية، كل مرحلة تُوزّع نوعاً إضافياً من حالات النموذج. التصميم مدروس: المراحل الأولى تُحقّق أكبر توفير في الذاكرة بأقل تكلفة تواصل.

افتح في المختبر
تنقّل بين مراحل ZeRO الثلاث وشاهد كيف تنخفض ذاكرة كل GPU مع كل مرحلة إضافية.
تستيقظ التجربة عند وصولك…

المرحلة الأولى — توزيع حالات المُحسِّن (P_os): كل GPU يحتفظ فقط بجزء 1/N من حالات المُحسِّن (حيث N عدد الأجهزة). بعد كل تمرير خلفي، تُجمع التدرّجات بالطريقة المعتادة عبر all-reduce، ثم يُحدّث كل جهاز حصّته فقط من المعاملات. بعد التحديث، توزّع المعاملات المُحدَّثة على الجميع. بهذا تنخفض ذاكرة المُحسِّن من 12Ψ إلى 12Ψ/N — توفير 4× على 64 جهازاً دون أي تواصل إضافي.

المرحلة الثانية — توزيع التدرّجات (P_os+g): بالبناء على المرحلة الأولى، كل GPU يحتفظ فقط بالتدرّجات المقابلة لحصّته من المُحسِّن. بدل all-reduce (التي تُعطي كل جهاز جميع التدرّجات)، يستخدم ZeRO (التي تُعطي كل جهاز شريحته فقط). بعد التحديث، تُوزَّع المعاملات كالسابق عبر all-gather. الذاكرة تنخفض من 14Ψ إلى (2 + 14/N)Ψ — توفير 8× على 64 جهازاً، وما زلنا دون تواصل إضافي.

المرحلة الثالثة — توزيع المعاملات (P_os+g+p): هنا يذهب ZeRO إلى أقصى حدّ، فكل GPU لا يحتفظ إلا بجزء 1/N من المعاملات ذاتها. أثناء والخلفي، تُستعاد معاملات كل طبقة كاملةً من الأجهزة الأخرى عبر all-gather، ثم تُحذف فور الانتهاء منها. الذاكرة تنخفض إلى 16Ψ/N — تقليص خطّي. الثمن هنا هو 1.5× من عبء التواصل (عمليتا all-gather إضافيتان في كل خطوة)، لكن هذا ما يفتح الباب لتدريب نماذج ضخمة حقاً.

Memory per GPU={4Ψ+12ΨNdStage 1 (Pos)2Ψ+14ΨNdStage 2 (Pos+g)16ΨNdStage 3 (Pos+g+p)\text{Memory per GPU} = \begin{cases} 4\Psi + \frac{12\Psi}{N_d} & \text{Stage 1 (P}_{os}\text{)} \\[6pt] 2\Psi + \frac{14\Psi}{N_d} & \text{Stage 2 (P}_{os+g}\text{)} \\[6pt] \frac{16\Psi}{N_d} & \text{Stage 3 (P}_{os+g+p}\text{)} \end{cases}
استهلاك الذاكرة لكل وحدة معالجة في ZeRO-DP — Ψ عدد معاملات النموذج، وN_d عدد أجهزة GPU المشاركة في التوازي. كل مرحلة تُوزّع نوعاً إضافياً من الحالات فتُخفّض ذاكرة كل جهاز أكثر. المرحلة الثالثة تمنحك تقليصاً خطّياً: إذا ضاعفتَ عدد الأجهزة انخفضت ذاكرة كل واحد للنصف.

ZeRO-R: ترويض الذاكرة المتبقّية

حتى بعد أن يُوزّع ZeRO-DP حالات النموذج، قد تبقى الذاكرة المتبقّية — التنشيطات والمخازن المؤقتة والتجزئة — هي العنق الضيّق. يتولّى ZeRO-R معالجة كل واحد منها.

توزيع التنشيطات والتفريغ إلى CPU: تقنية تُقلّل استهلاك الذاكرة بإعادة حساب التنشيطات أثناء التمرير الخلفي بدل تخزينها كلها. لكن ZeRO-R يخطو أبعد: يُوزّع التنشيطات المحفوظة بين الأجهزة، وفي النماذج الضخمة جداً يُفرّغها إلى ذاكرة CPU. ولأن التنشيطات لا تُستخدم إلا أثناء التمرير الخلفي، يمكن استرجاعها مسبقاً إلى GPU قبل الحاجة إليها بلحظات.

مخازن مؤقتة ثابتة الحجم: عمليات all-reduce تكون أكفأ حين تُدمج عدة مصفوفات صغيرة في مخزن واحد كبير، لكن تخصيص مخزن يتناسب حجمه مع حجم النموذج يخلق ضغطاً إضافياً على الذاكرة. الحلّ في ZeRO-R هو مخزن ثابت الحجم: كبير بما يكفي للكفاءة لكنه لا يكبر مع النموذج.

إلغاء تجزئة الذاكرة: أثناء التدريب، تُحجز الذاكرة وتُحرَّر بأحجام مختلفة فتتراكم فجوات لا يمكن استخدامها رغم أن المجموع الحرّ يبدو كافياً. يتجنّب ZeRO-R هذه المشكلة بحجز كتل ذاكرة متجاورة مسبقاً للتنشيطات والتدرّجات، فيضمن أن تبقى الذاكرة قابلة للاستخدام الفعلي.

افتح في المختبر
شاهد كيف يتعامل ZeRO-R مع كل مكوّن من مكوّنات الذاكرة المتبقّية.
تستيقظ التجربة عند وصولك…

تحليل التواصل: ثمن التقسيم

توزيع الذاكرة لا قيمة له إن كانت تكلفة التواصل بين الأجهزة باهظة. نجاح ZeRO يعتمد على اختيار ذكي لعمليات التواصل الجماعي.

في التوازي المعياري في البيانات، بعد التمرير الخلفي، تُزامن عملية all-reduce التدرّجات بين جميع الأجهزة. حجم البيانات المنقولة هو 2Ψ لكل GPU — هذا هو خط الأساس.

المرحلة الأولى تستبدل all-reduce بعملية reduce-scatter (كل جهاز يحصل على حصّته فقط من التدرّجات المجمَّعة) ثم all-gather (لتوزيع المعاملات المُحدَّثة). الحجم الإجمالي مطابق لخط الأساس — أي عبء صفري.

المرحلة الثانية تستخدم reduce-scatter بدل all-reduce، وهي عملياً أخفّ. مع all-gather بعد التحديث، يظلّ الحجم الكلّي مساوياً للأساس.

المرحلة الثالثة تُضيف عمليتَي all-gather — واحدة قبل التمرير الأمامي وأخرى قبل الخلفي — لإعادة تجميع المعاملات الكاملة من شرائحها. هذا يرفع التواصل إلى نحو 1.5× من الأساس. مع وصلات عالية مثل NVLink، هذه التكلفة مقبولة تماماً مقابل التوفير الهائل في الذاكرة.

افتح في المختبر
قارن أنماط التواصل بين أسلوب التوازي المعياري ومراحل ZeRO المختلفة.
تستيقظ التجربة عند وصولك…

التوسُّع إلى 100 مليار وما بعدها

أظهرت ورقة ZeRO نتائج توسُّع لافتة. باستخدام ZeRO-100B — الذي يجمع المرحلة الأولى مع التوازي في النموذج — درّب الباحثون نماذج تتجاوز 100 مليار معامل على 400 GPU من طراز NVIDIA V100، وحقّقوا تسارعاً فائق الخطّية و مستدامة تبلغ 15 بيتافلوبس.

الأهمّ من ذلك أن ZeRO أتاح التدريب واسع النطاق دون الحاجة إلى التوازي في النموذج أصلاً. باستخدام ZeRO-DP وحده، أمكن تدريب نماذج تصل إلى 13 مليار معامل — عشرة أضعاف ما يدعمه التوازي المعياري في البيانات. هذا أمر جوهري لأن التوازي في النموذج يتطلّب كتابة شيفرة مخصّصة لكل بنية على حدة، بينما ZeRO يعمل كبديل مباشر لا يحتاج تعديل شيفرة التدريب.

التحليل النظري في الورقة أظهر أنه مع 1024 GPU، تستطيع المرحلة الثالثة دعم نماذج بأكثر من تريليون معامل — نحو 6 أضعاف ما يتحمّله التوازي في النموذج وحده على نفس العتاد.

افتح في المختبر
احسب حجم النموذج الذي تستطيع تدريبه على عنقود GPU لديك مع كل مرحلة من ZeRO.
تستيقظ التجربة عند وصولك…

ZeRO عملياً: إعداد DeepSpeed

تفعيل المرحلة الثانية من ZeRO في DeepSpeedpython

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

import deepspeed

# إعداد DeepSpeed مع المرحلة الثانية من ZeRO
ds_config = {
    "train_batch_size": 32,
    "fp16": {"enabled": True},        # التدريب بالدقة المختلطة
    "zero_optimization": {
        "stage": 2,                    # تقسيم حالات المُحسِّن + التدرّجات
        "allgather_partitions": True,   # جمع المعاملات المُحدَّثة بعد خطوة المُحسِّن
        "reduce_scatter": True,         # استخدام reduce-scatter بدل all-reduce
        "overlap_comm": True,           # تداخل التواصل مع الحوسبة
        "contiguous_gradients": True,   # تقليل تجزئة الذاكرة
    },
}

# تغليف النموذج والمُحسِّن — ZeRO شفّاف لحلقة التدريب
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config,
)

# حلقة التدريب تبقى مطابقة لـ PyTorch المعياري
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

عائلة ZeRO وما مكّنته

  1. 2019

    ZeRO (الورقة الأصلية)

    توزيع حالات المُحسِّن والتدرّجات والمعاملات عبر ثلاث مراحل تصاعدية. تدريب نماذج بأكثر من 100 مليار معامل على 400 GPU بتسارع فائق الخطّية. نُشرت في SC 2020.

  2. 2020

    DeepSpeed وTuring-NLG

    أُدمج ZeRO ضمن منصة DeepSpeed، واستُخدم لتدريب Turing-NLG بـ 17.2 مليار معامل — أكبر نموذج لغوي حينها — محقّقاً دقة قياسية.

  3. 2021

    ZeRO-Offload

    امتداد لـ ZeRO يُفرّغ حالات المُحسِّن والتدرّجات إلى ذاكرة CPU، فأصبح ممكناً تدريب نماذج بمليارات المعاملات على GPU واحد. فتح الباب لتدريب النماذج الكبيرة بموارد محدودة.

  4. 2021

    ZeRO-Infinity

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

  5. 2022

    PyTorch FSDP

    تبنّت PyTorch أفكار ZeRO الأساسية في FSDP (التوازي الكامل المُجزَّأ في البيانات)، فأصبح أسلوب ZeRO في التدريب مدمجاً بشكل أصيل في منظومة PyTorch.

  6. 2023

    ZeRO++

    أضاف ضغط التواصل عبر التكميم وتقسيماً هرمياً يُقلّل حركة البيانات بين العقد حتى 4×، ممّا وسّع كفاءة ZeRO لتشمل العناقيد ذات الروابط البطيئة.

أثر ZeRO يتجاوز حدود ورقة بحثية أو نظام بعينه. ما أثبته هو أن جدار الذاكرة في تدريب ليس قيداً جوهرياً — بل نتيجة لتكرار البيانات دون حاجة. حين تُعامَل حالات النموذج كموارد مشتركة يمكن توزيعها وجمعها وتحريرها حسب الطلب، تتحوّل مشكلة الذاكرة من ضربية (كل GPU يحمل كل شيء) إلى جمعية (كل جهاز يحمل جزءاً فقط). هذا المبدأ يقف اليوم وراء كل نظام تدريب حديث واسع النطاق، من في Megatron-LM إلى ، ولا يزال الأساس الذي يُبنى عليه تدريب النماذج بتريليون معامل.

المرجعRajbhandari, Rasley, Ruwase, He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC (International Conference for High Performance Computing, Networking, Storage and Analysis), 2020.

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