التوازي الكامل للبيانات المجزأة

لتسريع تدريب النماذج الضخمة على أحجام دفعات أكبر، يمكننا استخدام نموذج موازٍ للبيانات مُجزّأ بالكامل. تُمكّن هذه البنية الموازية للبيانات من استيعاب المزيد من البيانات ونماذج أكبر من خلال تجزئة حالات المُحسِّن والتدرجات والمعاملات. للقراءة المزيد حول ذلك وفوائده، راجع مدونة FSDP. لقد قمنا بدمج أحدث ميزة تدريب FSDP في PyTorch. كل ما عليك فعله هو تفعيلها من خلال ملف الإعداد.

كيف يعمل دون أي إعداد مسبق

على جهازك أو أجهزتك، ما عليك سوى تشغيل:

accelerate config

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

accelerate launch my_script.py --args_to_my_script

على سبيل المثال، إليك كيفية تشغيل examples/nlp_example.py (من جذر المستودع) مع تفعيل FSDP:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

accelerate launch examples/nlp_example.py

حاليًا، يدعم Accelerate الإعدادات التالية عبر سطر الأوامر:

fsdp_sharding_strategy: [1] FULL_SHARD (يقسّم حالات المُحسِّن والتدرّجات والمعاملات)، [2] SHARD_GRAD_OP (يقسّم حالات المُحسِّن والتدرّجات)، [3] NO_SHARD (DDP)، [4] HYBRID_SHARD (يقسّم حالات المُحسِّن والتدرّجات والمعاملات داخل كل عقدة بينما تمتلك كل عقدة نسخة كاملة)، [5] HYBRID_SHARD_ZERO2 (يقسّم حالات المُحسِّن والتدرّجات داخل كل عقدة بينما تمتلك كل عقدة نسخة كاملة). لمزيد من المعلومات، يُرجى الرجوع إلى توثيقات PyTorch الرسمية.

fsdp_offload_params : يحدد ما إذا كان سيتم نقل المعاملات والتدرجات إلى وحدة المعالجة المركزية (CPU)

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP، [2] SIZE_BASED_WRAP، [3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: ينطبق على Transformers فقط. عند استخدام fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP، يمكن للمستخدم تقديم سلسلة مفصولة بفواصل لأسماء فئات طبقات الترانسفورمر (مع مراعاة حالة الأحرف) التي يجب تغليفها، مثل BertLayer، GPTJBlock، T5Block، BertLayer,BertEmbeddings,BertSelfOutput. هذا أمر مهم لأن الوحدات الفرعية التي تتشارك الأوزان (مثل طبقات التضمين) يجب ألا تنتهي في وحدات FSDP المغلفة المختلفة. باستخدام هذه السياسة، يتم التغليف لكل كتلة تحتوي على Multi-Head Attention تليها عدة طبقات MLP. تُغلَّف الطبقات المتبقية بما في ذلك التضمينات المشتركة بشكل مريح في وحدة FSDP الخارجية الأوسع. لذلك، استخدم هذا للنماذج القائمة على الترانسفورمر. يمكنك استخدام model._no_split_modules لنماذج Transformer عن طريق الإجابة بـ yes على Do you want to use the model's_no_split_modulesto wrap. It will try to usemodel._no_split_modules` when possible.

fsdp_min_num_params: الحد الأدنى لعدد المعاملات عند استخدام fsdp_auto_wrap_policy=SIZE_BASED_WRAP.

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE، [2] BACKWARD_POST، [3] NO_PREFETCH

fsdp_forward_prefetch: إذا كانت القيمة True، فإن FSDP يقوم مسبقًا بجلب (prefetch) عملية all-gather التالية أثناء تنفيذ مسار التقدّم (forward pass). يجب استخدام هذا الخيار فقط مع نماذج ذات رسم بياني ثابت (static-graph) لأن عملية الجلب المسبق تتبع ترتيب التنفيذ الخاص بالتكرار الأول. أي، إذا تغيّر ترتيب الوحدات الفرعية ديناميكيًا أثناء تنفيذ النموذج، فلا تمكّن هذه الميزة.

fsdp_state_dict_type: [1] FULL_STATE_DICT، [2] LOCAL_STATE_DICT، [3] SHARDED_STATE_DICT

fsdp_use_orig_params: إذا كانت القيمة True، فإنها تسمح بـ requires_grad غير منتظم أثناء التهيئة، مما يعني دعم المعاملات المجمدة والقابلة للتدريب المتناوبة. هذا الإعداد مفيد في حالات مثل الضبط الدقيق الفعال للمعاملات كما هو موضح في هذا المنشور. كما يتيح هذا الخيار وجود مجموعات معاملات متعددة للمحسّن. يجب أن تكون القيمة True عند إنشاء المحسّن قبل تجهيز/تغليف النموذج باستخدام FSDP.

fsdp_cpu_ram_efficient_loading: ينطبق هذا الخيار على نماذج Transformers فقط. إذا كان True، فإن العملية الأولى فقط هي التي تُحمّل نقطة التفتيش للنموذج المُدرَّب مسبقًا، بينما تكون الأوزان فارغة في جميع العمليات الأخرى. يجب ضبط هذا الخيار على False إذا واجهت أخطاء عند تحميل نموذج Transformers المُدرَّب مسبقًا عبر طريقة from_pretrained. عندما يكون هذا الإعداد True، يجب أيضًا أن يكون fsdp_sync_module_states True، وإلا فإن جميع العمليات باستثناء العملية الرئيسية ستحتوي على أوزان عشوائية مما يؤدي إلى سلوك غير متوقع أثناء التدريب. لكي يعمل هذا، تأكد من تهيئة مجموعة العمليات الموزعة قبل استدعاء طريقة from_pretrained في Transformers. عند استخدام واجهة Trainer API، يتم تهيئة مجموعة العمليات الموزعة عند إنشاء كائن من فئة TrainingArguments.

fsdp_sync_module_states: إذا كانت القيمة True، فسيبث كل وحدة FSDP ملفوفة بشكل فردي معاملات الوحدة من الرتبة 0.

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

فيما يلي مثال:

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

الحفظ والتحميل

الطريقة الجديدة الموصى بها لإنشاء نقاط التفتيش عند استخدام نماذج FSDP هي استخدام SHARDED_STATE_DICT كـ StateDictType عند إعداد تكوين accelerate. فيما يلي مقتطف الكود للحفظ باستخدام أداة save_state في accelerate.

accelerator.save_state("ckpt")

افحص مجلد نقطة التفتيش لرؤية النموذج والمُحسِّن على شكل أجزاء لكل عملية:

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

لتحميلها مرة أخرى لاستئناف التدريب، استخدم أداة load_state في accelerate

accelerator.load_state("ckpt")

عند استخدام save_pretrained في transformers، مرّر state_dict=accelerator.get_state_dict(model) لحفظ state dict للنموذج. فيما يلي مثال:

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

State Dict

ستقوم accelerator.get_state_dict باستدعاء تنفيذ model.state_dict الأساسي باستخدام مدير السياق FullStateDictConfig(offload_to_cpu=True, rank0_only=True) للحصول على قاموس الحالة (state dict) للرتبة 0 فقط، وسيتم نقله إلى وحدة المعالجة المركزية (CPU).

يمكنك بعد ذلك تمرير state إلى طريقة save_pretrained. هناك عدة أوضاع لـ StateDictType و FullStateDictConfig يمكنك استخدامها للتحكم في سلوك state_dict. لمزيد من المعلومات، راجع توثيق PyTorch.

إذا اخترت استخدام StateDictType.SHARDED_STATE_DICT، فسيتم تقسيم أوزان النموذج أثناء Accelerator.save_state إلى n ملفات لكل تقسيم فرعي على النموذج. لدمجها مرة أخرى في قاموس واحد لتحميلها لاحقًا في النموذج بعد التدريب، يمكنك استخدام أداة merge_weights:

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

سيتم حفظ المخرجات النهائية إما في model.safetensors أو pytorch_model.bin (إذا تم تمرير safe_serialization=False).

يمكن أيضًا استدعاء هذا الأمر باستخدام سطر الأوامر:

accelerate merge-weights pytorch_model_fsdp_0/ output_path

التطابق بين استراتيجيات تقسيم FSDP ومراحل DeepSpeed ZeRO

• يتطابق FULL_SHARD مع ZeRO Stage-3 في DeepSpeed. يقوم بتجزئة حالات المُحسِّن والتدرجات والمعاملات.
• يتطابق SHARD_GRAD_OP مع ZeRO Stage-2 في DeepSpeed. يقوم بتجزئة حالات المُحسِّن والتدرجات.
• يتطابق NO_SHARD مع ZeRO Stage-0. لا يتم تقسيم النموذج، حيث تمتلك كل وحدة معالجة رسومية نسخة كاملة من النموذج وحالات المُحسِّن والتدرجات.
• يُشير HYBRID_SHARD إلى ZeRO++ Stage-3 حيث يكون zero_hpz_partition_size=<num_gpus_per_node>. هنا، يتم تقسيم حالات المُحسِّن والتدرجات والمعاملات داخل كل عقدة، بينما تمتلك كل عقدة نسخة كاملة.

بعض التحذيرات التي يجب الانتباه إليها

• في حالة وجود عدة نماذج، قم بتمرير المُحسِّنين إلى استدعاء prepare بنفس ترتيب النماذج المقابلة، وإلا فإن accelerator.save_state() وaccelerator.load_state() ستؤدي إلى سلوك خاطئ أو غير متوقع.
• هذا الميزة غير متوافقة مع --predict_with_generate في سكربت run_translation.py التابع لمكتبة Transformers.

للحصول على مزيد من التحكم، يمكن للمستخدمين الاستفادة من FullyShardedDataParallelPlugin. بعد إنشاء نسخة من هذه الفئة، يمكن للمستخدمين تمريرها إلى إنشاء نسخة من فئة Accelerator. لمزيد من المعلومات حول هذه الخيارات، يُرجى الرجوع إلى كود FullyShardedDataParallel الخاص بـ PyTorch.

للمهتمين بالتشابهات والاختلافات بين FSDP وDeepSpeed، يرجى الاطلاع على دليل المفاهيم هنا!