DeepSpeed ZeRO

يُزيل DeepSpeed ZeRO (Zero Redundancy Optimizer) التكرار في الذاكرة عبر التدريب الموزّع من خلال تقسيم حالات المُحسِّن والتدرّجات والمعاملات على وحدات معالجة الرسوميات. يتكوّن ZeRO من ثلاث مراحل، حيث تقسم كل مرحلة حالة أكثر من سابقتها. كما يدعم DeepSpeed التفريغ إلى ذاكرة المعالج المركزي أو ذاكرة NVMe لتحقيق توفير إضافي. تقلّل كل مرحلة إضافية ومستوى تفريغ من الذروة في استهلاك الذاكرة، على حساب زيادة التواصل بين وحدات معالجة الرسوميات.

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

يقوم ZeRO-2 بتجزئة التدرجات وحالات المُحسِّن مع انخفاض في تكلفة الاتصال مقارنةً بـ ZeRO-3. استخدم ZeRO-3 فقط عندما لا يتسع نموذجك عبر وحدات معالجة الرسومات باستخدام ZeRO-2.

التثبيت

قم بتثبيت DeepSpeed من PyPI، أو ثبّت Transformers مع الإضافة deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

إذا واجهت أخطاء تثبيت متعلقة بـ CUDA، فراجع وثائق DeepSpeed CUDA. يُعد التثبيت من المصدر الخيار الأكثر موثوقية لأنه يتوافق مع عتادك الدقيق ويتضمن ميزات غير متاحة بعد في إصدار PyPI.

إعداد

يتكامل Trainer مع DeepSpeed من خلال الحجة deepspeed، والتي تقبل ملف إعداد بصيغة JSON. أو استخدم ملف إعداد Accelerate بدلاً من TrainingArguments.

TrainingArguments

ملف إعدادات Accelerate

استخدم "auto" في إعداداتك للقيم التي تريد أن يملأها DeepSpeed من TrainingArguments. إذا كنت تريد تحديد قيمة بشكل صريح، فتأكد من استخدام نفس القيمة لكل من معامل DeepSpeed وTrainingArguments.

راجع مرجع DeepSpeed Configuration JSON للحصول على قائمة كاملة بخيارات تكوين DeepSpeed.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

إعدادات DeepSpeed	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

مرّر الإعداد إلى معامل deepspeed.

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

مراحل ZeRO

اختر إعداد مرحلة ZeRO لاستخدامه كنقطة انطلاق.

ZeRO-1

ZeRO-2

ZeRO-3

الحقول التالية مهمة لتخصيص التدريب.

zero_optimization يحدد مرحلة ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
اضبط حجج حجم الدفعة وتراكم التدرجات على "auto". إذا قمت بإعداد هذه القيم يدوياً بقيم لا تتوافق مع TrainingArguments، فسيستمر التدريب بصمت بالقيم الخاطئة.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 يحدد دقة التدريب. اضبطه على "auto" ليتطابق مع علم bf16 في TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save يُنفّذ عملية all-gather عبر جميع وحدات معالجة الرسوميات قبل الحفظ، لإعادة بناء الموترات الكاملة من مقاطعها. هذا هو حجة ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
اضبط overlap_comm على true لإخفاء زمن تأخير all-reduce خلف مسار الرجوع. وتُوازن allgather_bucket_size وreduce_bucket_size بين سرعة الاتصال وذاكرة وحدة معالجة الرسوميات. فالقيم الأقل تستخدم ذاكرة أقل لكنها تُبطئ الاتصال.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer ينقل المُحسِّن إلى ذاكرة المعالج المركزي. ولتوفير المزيد من الذاكرة، انقل أيضًا معاملات النموذج باستخدام offload_param (لنسخة ZeRO-3 فقط). اضبط pin_memory على true لتسريع عمليات النقل بين المعالج المركزي والوحدة الرسومية، لكن هذا يُثبّت ذاكرة الوصول العشوائي فلا تصبح متاحةً للعمليات الأخرى.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```
optimizer وscheduler افتراضيًا هما المُحسِّن والمُجدوِل المُعيَّنان في TrainingArguments. اضبطهما على "auto" حتى يقرأ DeepSpeed القيم من TrainingArguments ما لم تكن بحاجة إلى مُحسِّن أصلي لـ DeepSpeed مثل LAMB.
```
{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}
```
إذا كنت تقوم بإسناد المُحسِّن إلى وحدة المعالجة المركزية، فاضبط zero_force_ds_cpu_optimizer على false لاستخدام مُحسِّن Adam على وحدة المعالجة المركزية من DeepSpeed.
```
{
    "zero_force_ds_cpu_optimizer": false
}
```

نقاط التفتيش

يحفظ DeepSpeed نقاط التفتيش في تنسيق مجزأ لا يمكن تحميله مباشرة باستخدام from_pretrained(). اضبط load_best_model_at_end() على True ليتمكن Trainer من تتبع أفضل نقطة تفتيش وإعادة تحميلها في نهاية التدريب.

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

تؤدي تعيين save_only_model=True إلى تخطي حفظ حالة المُحسِّن الكاملة، مما يعني أنك لن تتمكن من إعادة تحميل أفضل نموذج في نهاية التدريب. كما اضبط stage3_gather_16bit_weights_on_model_save: true لإعادة بناء الأوزان الكاملة من شقوقها. هذا مطلوب لحفظ نموذج موحّد بدقة 16 بت أو حالة dict بدقة 16 بت مع ZeRO-3. يرمي Transformers خطأً عند دمج save_only_model=True مع load_best_model_at_end=True.

لمتابعة الاستئناف عبر تكوينات التوازي المختلفة، راجع دليل Universal Checkpointing الخاص بـ DeepSpeed.

الخطوات التالية

اقرأ الفصل Zero Redundancy Optimizer من The Ultra-Scale Playbook لمعرفة المزيد حول كيفية عمل ZeRO.
اقرأ أبحاث ZeRO: Memory Optimizations Toward Training Trillion Parameter Models، وDemocratizing Billion-Scale Model Training، وBreaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub