FSDP2

يقوم Fully Sharded Data Parallel (FSDP2) بتقسيم النموذج والتدرجات وحالات المُحسِّن عبر وحدات معالجة الرسوميات. قبل الحساب، تجمع كل وحدة معالجة رسوميات مجموعة كاملة من المعلمات من جميع الأجزاء، ثم تفرّغها بعد ذلك. يتيح لك التقسيم تدريب نماذج أكبر من ذاكرة وحدة معالجة رسوميات واحدة، مقابل زيادة في حجم الاتصالات مقارنةً بـ DDP. استخدم FSDP عندما لا يتسع نموذجك أو حالات المُحسِّن على وحدة معالجة رسوميات واحدة.

                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step

استراتيجيات التجزئة

يتحكم FSDP2 في تقسيم النموذج عبر ~TrainingArguments.fsdp_config. اضبط fsdp=True لتفعيل FSDP، واضبط reshard_after_forward في إعدادات FSDP لاختيار التوازن بين الذاكرة ومعدل الإنتاجية.

يتحكم auto_wrap_policy في كيفية تغليف الوحدات داخل وحدات FSDP. القيمة الافتراضية هي "TRANSFORMER_BASED_WRAP"، والتي تقوم بتغليف طبقات الترانسفورمر في النموذج. بدون التغليف ("NO_WRAP")، يصبح النموذج بأكمله وحدة FSDP واحدة، فتخسر بذلك فائدة تقسيم الذاكرة.

إعداد FSDP

تحكم هذه الحقول في كيفية تغليف FSDP2 للنموذج وتجزئته وتحميله. يُغطي استراتيجيات التجزئة كلًا من reshard_after_forward وauto_wrap_policy.

• cpu_offload ينقل المعاملات والتدرجات إلى وحدة المعالجة المركزية (CPU) عندما لا تكون قيد الاستخدام، وذلك لتوفير ذاكرة وحدة معالجة الرسوميات (GPU).

• transformer_layer_cls_to_wrap يحدد طبقة الترانسفورمر التي سيتم تغليفها في وحدة FSDP عندما يكون auto_wrap_policy هو "TRANSFORMER_BASED_WRAP". تدير كل وحدة عمليات الجمع والتشتت الخاصة بها. يتم جمع معاملات الوحدة الحالية فقط أثناء المرور الأمامي. يتم تحرير معاملات الوحدات السابقة لتوفير الذاكرة.

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

• min_num_params يحدد الحد الأدنى لعدد المعاملات لكل وحدة للتغليف القائم على الحجم. يُستخدم فقط عندما يكون auto_wrap_policy هو "SIZE_BASED_WRAP".

• يتحكم state_dict_type في صيغة نقطة الحفظ. القيمة الافتراضية هي "FULL_STATE_DICT" لنقطة حفظ واحدة متوافقة مع Transformers. استخدم "SHARDED_STATE_DICT" للحصول على ملف نقطة حفظ واحد لكل rank، وهو أسرع للنماذج الكبيرة. لا يمكن تحميل نقاط الحفظ المجزأة إلا في FSDP، لذا احفظ "FULL_STATE_DICT" لنقطة الحفظ النهائية التي تريد مشاركتها أو تحميلها خارج FSDP.

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

• يعيد activation_checkpointing حساب التفعيلات أثناء المرور العكسي بدلاً من تخزينها. استخدم هذا بدلاً من gradient checkpointing في TrainingArguments. يؤدي تعيين كلاهما إلى رفع خطأ.

قم بإعداد تدريب FSDP إما باستخدام ملف إعدادات Accelerate أو ملف إعدادات FSDP يُمرَّر إلى fsdp_config.

accelerate launch train.py

{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

في نماذج مزيج الخبراء، يمكن دمج fsdp_size مع tp_size وep_size لتقسيم الخبراء عبر بُعد واحد من الشبكة وكل ما عداها عبر البُعد الآخر. راجع التوازي بين الخبراء.

التقييم باستخدام توازٍ السياق

يتم ضبط التوازي السياقي (CP) مع FSDP2 عبر ParallelismConfig في Accelerate باستخدام cp_size. عند الإقلاع، مرّر كلًا من --use_parallelism_config و--parallelism_config_cp_size. كما أن وجود parallelism_config غير فارغ في ملف YAML الخاص بـ Accelerate يُفعّل use_parallelism_config أيضًا. أما تمرير --parallelism_config_cp_size بمفرده فيُتجاهل ويبقي على cp_size=1.

في إطار FSDP2، يقوم CP بتجزئة التسلسل عبر الرتب بحيث تحتفظ كل وحدة معالجة رسومية (GPU) بشريحة من السياق (allgather أو alltoall). يختلف ذلك عن التوازي التسلسلي Ulysses، الذي يعمل عبر DeepSpeed ALST ويستخدم sp_size في نفس الإعداد. اختر بين CP أو Ulysses SP. لا تضبط كلًا من cp_size وsp_size على قيمة أكبر من 1.

يستخدم evaluate() نفس سياق التوازي السياقي الذي يستخدمه train() عبر prediction_step(). يقوم كل رتبة بتشغيل شريحة تسلسله بينما يكون ذلك السياق نشطاً. eval_loss هو الخسارة الصحيحة (غير مضروبة في cp_size). يستخدم التوقف المبكر وload_best_model_at_end ذلك المقياس.

تبقى القيم اللوجستية (Logits) مجزأة على طول التسلسل، ولا يوجد ما يعيد تجميعها إلى تنبؤات بطول كامل. لا يتم دعم التنبؤات وcompute_metrics على القيم اللوجستية في وضع CP.

أعد استخدام عملية الإقلاع الحالية الخاصة بك بـ FSDP2 + cp_size عندما تحتاج إلى eval_loss.

إذا كنت تُعرّف فئة فرعية من prediction_step() وتقوم بعدّ الرموز بنفسك، فعدّ num_items_in_batch قبل الدخول إلى سياق التوازي السياقي. بينما يكون هذا السياق نشطًا، تُقسَّم مخازن التسلسل في مكانها عبر الرتب. العدّ بعد ذلك يؤدي إلى نقص في عدد عناصر الدفعة. الفئة الافتراضية Trainer تعدّ أولاً بالفعل، ولا تحتاج إلى هذا إلا إذا تجاوزت prediction_step.

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

• راجع DDP للتدريب المتوازي على مستوى البيانات عندما يتسع نموذجك على وحدة معالجة رسومية واحدة.
• راجع DeepSpeed لتحسين ZeRO وإزاحة NVMe.
• لـ FSDP على TPUs مع PyTorch/XLA، اضبط xla وxla_fsdp_settings وxla_fsdp_grad_ckpt في ~TrainingArguments.fsdp_config.
• Read the FSDP chapter from The Ultra-Scale Playbook for more information about how FSDP works.