تحميل النماذج الكبيرة في الذاكرة

عند تحميل نموذج مُدرَّب مسبقًا في PyTorch، يبدو سير العمل المعتاد على النحو التالي:

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

ببساطة، تلك الخطوات هي:

1. أنشئ النموذج بأوزان مهيأة عشوائياً
2. حمّل أوزان النموذج (في قاموس يُسمّى عادةً state dict) من القرص
3. حمّل تلك الأوزان داخل النموذج

بينما يعمل هذا بشكل ممتاز مع النماذج ذات الأحجام المعتادة، فإن لهذا سير العمل بعض القيود الواضحة عند التعامل مع نموذج ضخم: في الخطوة 1، نقوم بتحميل نسخة كاملة من النموذج في الذاكرة العشوائية (RAM)، ونقضي بعض الوقت في تهيئة الأوزان عشوائياً (والتي سيتم التخلص منها في الخطوة 3). في الخطوة 2، نقوم بتحميل نسخة كاملة أخرى من النموذج في الذاكرة العشوائية، مع الأوزان المُدرَّبة مسبقاً. إذا كنت تقوم بتحميل نموذج يحتوي على 6 مليارات من المعاملات، فهذا يعني أنك ستحتاج إلى 24 جيجابايت من الذاكرة العشوائية لكل نسخة من النموذج، أي 48 جيجابايت إجمالاً (نصفها لتحميل النموذج بدقة FP16).

هذا الواجهة البرمجية (API) جديدة إلى حد كبير ولا تزال في مرحلتها التجريبية. وعلى الرغم من أننا نسعى جاهدين لتقديم واجهة برمجية مستقرة، إلا أنه من الممكن أن تتغير بعض الأجزاء الصغيرة من الواجهة البرمجية العامة في المستقبل.

كيف يعمل العملية: نظرة عامة سريعة

كيف يعمل العملية: العمل مع الكود

إنشاء نموذج فارغ

أول أداة يقدمها Accelerate للمساعدة في التعامل مع النماذج الكبيرة هي مدير سياق init_empty_weights() يساعدك على تهيئة نموذج دون استخدام أي ذاكرة RAM، بحيث يمكن تنفيذ الخطوة 1 على نماذج بأي حجم. وإليك كيف يعمل:

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

على سبيل المثال:

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

يُهيّئ نموذجًا فارغًا يحتوي على ما يزيد قليلاً عن 100 مليار معامل. في الخلفية، يعتمد هذا على جهاز meta الذي أُدخل في PyTorch 1.9. أثناء التهيئة ضمن مدير السياق، يتم نقل كل معامل فور إنشائه إلى ذلك الجهاز.

لا يمكنك نقل نموذج مُهيّأ بهذه الطريقة على وحدة المعالجة المركزية (CPU) أو جهاز آخر مباشرةً، لأنه لا يحتوي على أي بيانات. كما أن من المرجح جدًا أن تفشل عملية التمرير الأمامي (forward pass) باستخدام ذلك النموذج الفارغ، حيث لا تدعم جميع العمليات جهاز الميتا (meta device).

نقاط التفتيش المجزأة

من الممكن أن يكون حجم نموذجك كبيراً لدرجة أن نسخة واحدة منه لا تتسع في الذاكرة العشوائية (RAM). وهذا لا يعني أنه لا يمكن تحميله: إذا كان لديك جهاز أو عدة أجهزة معالجة رسومية (GPU)، فهذا يعني توفر ذاكرة إضافية لتخزين نموذجك. في هذه الحالة، من الأفضل أن يكون نقطة الفحص (checkpoint) مقسمة إلى عدة ملفات أصغر نسميها شظايا نقطة الفحص (checkpoint shards).

ستتعامل Accelerate مع نقاط التفتيش المجزأة ما دمتَ تتبع التنسيق التالي: يجب أن يكون نقطة التفتيش في مجلد، مع عدة ملفات تحتوي على قواميس الحالة الجزئية، ويجب أن يكون هناك فهرس بتنسيق JSON يحتوي على قاموس يربط أسماء المعاملات بالملف الذي يحتوي على أوزانها. يمكنك بسهولة تجزئة نموذجك باستخدام save_model(). على سبيل المثال، يمكن أن يكون لدينا مجلد يحتوي على:

first_state_dict.bin
index.json
second_state_dict.bin

حيث يكون ملف index.json هو الملف التالي:

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

وfirst_state_dict.bin الذي يحتوي على الأوزان الخاصة بـ"linear1.weight" و"linear1.bias"، وsecond_state_dict.bin الذي يحتوي على الأوزان الخاصة بـ"linear2.weight" و"linear2.bias"

تحميل الأوزان

الأداة الثانية التي يقدمها Accelerate هي الدالة load_checkpoint_and_dispatch()، والتي تتيح لك تحميل نقطة تحقق داخل نموذجك الفارغ. تدعم هذه الدالة نقاط التحقق الكاملة (ملف واحد يحتوي على كامل state dict) وكذلك نقاط التحقق المقسمة. كما أنها توزّع هذه الأوزان تلقائيًا عبر الأجهزة المتاحة لديك (GPUs، ذاكرة CPU)، لذا إذا كنت تقوم بتحميل نقطة تحقق مقسمة، فإن أقصى استخدام للذاكرة سيكون حجم أكبر شريحة.

إذا كنت ترغب في استخدام الاستدلال على النماذج الكبيرة مع نماذج Transformers، فاطّلع على هذه التوثيق.

إليك كيفية استخدام ذلك لتحميل نموذج GPT2-1.5B.

لنقم بتنزيل النسخة المقسمة من هذا النموذج.

pip install huggingface_hub

from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

لتهيئة النموذج، سنستخدم مكتبة minGPT.

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

ثم، حمّل نقطة التفتيش التي قمنا بتنزيلها للتو باستخدام:

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

عبر تمرير device_map="auto"، نُملي على Accelerate تحديد تلقائيًا مكان وضع كل طبقة من طبقات النموذج اعتمادًا على الموارد المتاحة:

• أولاً، نستخدم الحد الأقصى من المساحة المتاحة على وحدة/وحدات معالجة الرسوميات (GPU)
• إذا كنا ما زلنا بحاجة إلى مساحة، نخزّن الأوزان المتبقية على وحدة المعالجة المركزية (CPU)
• إذا لم يكن هناك ما يكفي من ذاكرة الوصول العشوائي (RAM)، نخزّن الأوزان المتبقية على القرص الصلب على شكل متجهات مرمّزة بالذاكرة (memory-mapped tensors)

no_split_module_classes

ستُشير هذه المعلمة إلى أن بعض الوحدات التي تحمل الاسم "Block" يجب ألا تُقسَّم عبر أجهزة مختلفة. عليك أن تحدد هنا جميع الوحدات التي تتضمن اتصالاً متبقياً من نوع ما.

device_map

يمكنك الاطلاع على device_map الذي اختاره Accelerate عن طريق الوصول إلى خاصية hf_device_map في نموذجك:

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

من الممكن تمامًا إنشاء خريطة أجهزة خاصة بالطبقات لاستخدامها، مع تحديد جهاز GPU المطلوب (رقم)، أو "cpu"، أو "disk"، ثم تمرير هذه الخريطة:

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

تشغيل النموذج

الآن، بعد أن قمنا بذلك، يتوزّع نموذجنا على عدة أجهزة، وربما على القرص الصلب. لكن لا يزال بإمكاننا استخدامه كنموذج PyTorch عادي:

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

في الخلفية، أضافت Accelerate خطافات إلى النموذج، بحيث:

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

بهذه الطريقة، يمكن لنموذجك أن يعمل للاستدلال حتى لو لم يتسع على أحد وحدات معالجة الرسوميات أو ذاكرة الوصول العشوائي للمعالج!

هذا يدعم فقط الاستدلال على نموذجك، وليس التدريب. يحدث معظم الحساب خلف مُديري السياق torch.no_grad() لتجنب استهلاك جزء من ذاكرة وحدة معالجة الرسوميات (GPU) بواسطة التفعيلات الوسيطة.

تصميم خريطة جهاز

يمكنك ترك أمر حساب خريطة الأجهزة لـ Accelerate عن طريق تعيين device_map إلى أحد الخيارات المدعومة ("auto"، "balanced"، "balanced_low_0"، "sequential") أو إنشاء خريطة بنفسك إذا كنت ترغب في مزيد من التحكم في المكان الذي يجب أن تذهب إليه كل طبقة.

يمكنك اشتقاق جميع أحجام النموذج (وبالتالي حساب device_map) على نموذج موجود على جهاز meta.

ستنتج جميع الخيارات النتيجة نفسها عندما لا يتوفر لديك ذاكرة كافية على وحدة معالجة الرسوميات (GPU) لاستيعاب النموذج بالكامل (وهو ما يعني وضع كل ما يمكن وضعه على وحدة معالجة الرسوميات، ثم نقل الأوزان إلى وحدة المعالجة المركزية (CPU) أو حتى إلى القرص الصلب إذا لم تكن ذاكرة الوصول العشوائي (RAM) كافية).

عندما تكون ذاكرة وحدة معالجة الرسوميات المتاحة لديك أكبر من حجم النموذج، فإن الفرق بين كل خيار هو كالتالي:

• "auto" و"balanced" يقسمان النموذج بالتساوي على جميع وحدات معالجة الرسومات المتاحة، مما يتيح لك استخدام حجم دفعة أكبر من 1.
• يقوم "balanced_low_0" بتقسيم النموذج بالتساوي على جميع وحدات معالجة الرسومات باستثناء الوحدة الأولى، ولا يضع على وحدة GPU 0 سوى ما لا يتسع على الوحدات الأخرى. هذه الخيار ممتاز عندما تحتاج إلى استخدام وحدة GPU 0 لمعالجة بعض المخرجات، كما هو الحال عند استخدام دالة generate لنماذج Transformers
• "sequential" سيملأ ما يمكنه على وحدة معالجة الرسوميات 0، ثم ينتقل إلى وحدة معالجة الرسوميات 1 وهكذا (لذا لن يستخدم وحدات معالجة الرسوميات الأخيرة إذا لم يكن بحاجة إليها).

تنتج الخيارات "auto" و"balanced" نفس النتائج حاليًا، لكن سلوك "auto" قد يتغير في المستقبل إذا وجدنا استراتيجية أكثر منطقية، بينما سيبقى "balanced" مستقرًا.

لاحظ أولاً أنه يمكنك تحديد الذاكرة المستخدمة على كل وحدة معالجة رسومية (GPU) باستخدام الحجة max_memory (المتاحة في infer_auto_device_map() وفي جميع الدوال التي تستخدمها). عند تعيين max_memory، يجب أن تمرر قاموساً يحتوي على معرّفات وحدات المعالجة الرسومية (على سبيل المثال 0، 1 إلخ) ومفتاح "cpu" للحد الأقصى من ذاكرة الوصول العشوائي (RAM) التي تريد استخدامها لتفريغ المعالج المركزي (CPU offload). يمكن أن تكون القيم إما عددًا صحيحًا (بالبايت) أو نصاً يمثل رقماً مع وحدته، مثل "10GiB" أو "10GB".

فيما يلي مثال لا نرغب فيه باستخدام أكثر من 10GiB على كل من بطاقتي GPU ولا أكثر من 30GiB من ذاكرة CPU العشوائية (RAM) لأوزان النموذج:

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

عند حدوث أول تخصيص في PyTorch، يتم تحميل نواة CUDA التي تستهلك ما يقرب من 1-2 جيجابايت من الذاكرة اعتمادًا على وحدة معالجة الرسوميات. لذلك، تكون الذاكرة المتاحة دائمًا أقل من الحجم الفعلي لوحدة معالجة الرسوميات. لمعرفة مقدار الذاكرة المستخدمة فعليًا، قم بتنفيذ torch.ones(1).cuda() وراقب استخدام الذاكرة.

لذلك، عند إنشاء خرائط الذاكرة باستخدام max_memory، تأكد من ضبط الذاكرة المتاحة وفقاً لذلك لتجنب أخطاء نفاد الذاكرة.

بالإضافة إلى ذلك، إذا قمت ببعض العمليات الإضافية على مخرجاتك دون إعادتها إلى وحدة المعالجة المركزية (على سبيل المثال داخل طريقة generate في Transformers) وإذا وضعت مدخلاتك على وحدة معالجة رسومية، فإن وحدة المعالجة الرسومية هذه ستستهلك ذاكرة أكثر من الوحدات الأخرى (حيث يقوم Accelerate دائمًا بإعادة المخرجات إلى جهاز المدخلات). لذلك، إذا كنت ترغب في تحسين حجم الدفعة الأقصى ولديك وحدات معالجة رسومية متعددة، فامنح وحدة المعالجة الرسومية الأولى ذاكرة أقل. على سبيل المثال، مع BLOOM-176B على تكوين 8x80 A100، فإن الخريطة القريبة من المثالية هي:

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

كما ترون، منحنا الـ 7 وحدات معالجة الرسومات المتبقية ذاكرة أكبر بنسبة 50% تقريبًا مقارنةً بوحدة GPU 0.

إذا اخترت تصميم device_map بالكامل بنفسك، فيجب أن يكون قاموساً تكون مفاتيحه أسماء وحدات نموذجك وقيمها معرّف جهاز صالح (على سبيل المثال عدد صحيح للوحدات الرسومية) أو "cpu" لنقل المعالجة إلى المعالج، أو "disk" لنقل المعالجة إلى القرص. يجب أن تغطي المفاتيح النموذج بأكمله، ويمكنك بعد ذلك تحديد خريطة الأجهزة كما تشاء: على سبيل المثال، إذا كان نموذجك يحتوي على بلوكين (لنقل block1 وblock2) يحتوي كل منهما على ثلاث طبقات خطية (لنقل linear1، linear2 وlinear3)، فقد تكون خريطة الأجهزة الصالحة كالتالي:

device_map = {"block1": 0, "block2": 1}

مثال آخر صالح يمكن أن يكون:

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

من ناحية أخرى، هذا غير صالح لأنه لا يغطي كل معلمات النموذج:

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

لضمان أعلى كفاءة، تأكد من أن خريطة جهازك تضع المعاملات على وحدات معالجة الرسوميات (GPUs) بشكل متسلسل (على سبيل المثال، لا تضع أحد الأوزان الأولى على GPU 0، ثم الأوزان على GPU 1، ثم الوزن الأخير مرة أخرى على GPU 0) لتجنب إجراء العديد من عمليات نقل البيانات بين وحدات معالجة الرسوميات.

إزاحة المعالجة إلى وحدة المعالجة المركزية فقط

إذا كنت ترغب في نقل نموذجك إلى وحدة المعالجة المركزية (CPU)، يمكنك استخدام cpu_offload(). نتيجةً لذلك، سيتم نقل جميع معاملات النموذج، وسيتم الاحتفاظ بنسخة واحدة فقط من قاموس الحالة (state dict) للنموذج. أثناء المرور الأمامي (forward pass)، سيتم استخراج المعاملات من ذلك القاموس ووضعها على جهاز التنفيذ، ثم تمريرها عند الحاجة، ثم نقلها مرة أخرى.

cpu_offload(model, execution_device)

You can also use cpu_offload_with_hook(). This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with cpu_offload() is that the model stays on the execution device after the forward and is only offloaded again when the offload method of the returned hook is called. Furthermore, cpu_offload_with_hook() is more performant but less memory saving. It is useful for pipelines running a model in a loop:

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

إخلاء إلى القرص فقط

لتنفيذ إزاحة القرص، يمكنك استخدام disk_offload(). ونتيجةً لذلك، سيتم إزاحة جميع معاملات النموذج كمصفوفة مرسومة في الذاكرة داخل مجلد محدد. أثناء المرور الأمامي، سيتم الوصول إلى المعاملات من ذلك المجلد ووضعها على جهاز التنفيذ المُمرَّر كما هي عند الحاجة إليها، ثم تُزاح مرة أخرى.

disk_offload(model, offload_dir, execution_device)

القيود والتطوير الإضافي

نحن على دراية بالقيود الحالية في واجهة برمجة التطبيقات:

• يحاول infer_auto_device_map() (أو device_map="auto" في load_checkpoint_and_dispatch()) تعظيم ذاكرة GPU و CPU المتاحة التي يراها عند تنفيذك له. بينما يتقن PyTorch إدارة ذاكرة GPU بكفاءة (وإعادتها عند عدم الحاجة إليها)، فإن الأمر ليس كذلك تماماً بالنسبة إلى Python وذاكرة CPU. لذلك، قد يكون خريطة الأجهزة المحسوبة تلقائياً مرهقة جداً للمعالج. انقل بعض الوحدات إلى جهاز القرص إذا واجهت أعطالاً بسبب نقص الذاكرة.
• يعمل infer_auto_device_map() (أو device_map="auto" في load_checkpoint_and_dispatch()) على إسناد الأجهزة بشكل تسلسلي (لتجنب نقل البيانات ذهاباً وإياباً)، لذا إذا كانت طبقتك الأولى أكبر من حجم وحدة معالجة الرسوميات (GPU) المتوفرة لديك، فستنتهي العملية بوضع كل شيء على وحدة المعالجة المركزية (CPU) أو القرص.
• لا تُجري الدالتان load_checkpoint_and_dispatch() وload_checkpoint_in_model() أي فحص لصحة قاموس الحالة الخاص بك مقارنةً بنموذجك في الوقت الحالي (سيتم إصلاح ذلك في إصدار مستقبلي)، لذا قد تظهر لك بعض الأخطاء الغريبة عند محاولة تحميل نقطة تحقق تحتوي على مفاتيح غير متطابقة أو مفاتيح مفقودة.
• إن توازي النموذج المستخدم عند تقسيم نموذجك على عدة وحدات معالجة رسومية (GPU) هو توازي بسيط وغير مُحسَّن، مما يعني أن وحدة معالجة رسومية واحدة فقط تعمل في أي وقت معيّن بينما تبقى الوحدات الأخرى خاملة.
• عندما يتم تفريغ الأوزان على وحدة المعالجة المركزية/قرص التخزين الصلب، لا يوجد جلب مسبق (ولكننا سنعمل على ذلك في الإصدارات المستقبلية)، مما يعني أن الأوزان تُوضع على وحدة معالجة الرسوميات عندما تكون مطلوبة وليس قبل ذلك.
• قد يكون تفريغ البيانات إلى القرص الصلب بطيئًا للغاية إذا لم يكن لديك اتصال سريع بين القرص ومعالج وحدة المعالجة المركزية (مثل NVMes) على العتاد الذي تعمل عليه.