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

قد لا يهتم العديد من المهندسين أو الباحثين في هذا المجال بكيفية استهلاك ذاكرة LLM، وخاصةً في مرحلة التدريب. عندما بدأت العمل على LLM، كان من المحبط أن أحصل على خطأ OOM (نفاد الذاكرة) كلما حاولت تغيير بعض المعلمات باستخدام نص التدريب. لقد حاولت معرفة سبب حدوث ذلك لبنية المحولات، لكنني لم أجد أي مقال مفيد حول هذا الموضوع.

نحن نعلم أن معظم النماذج الحديثة في مجال معالجة اللغة الطبيعية هي نماذج فك تشفير المحولات فقط مثل GPT وLlama وMistral وPhi وما إلى ذلك. لذا قررت فهم استخدام الذاكرة لهذه الأنواع من النماذج باستخدام الحساب اليدوي. لقد أجريت حسابًا يدويًا حول كيفية استهلاك الذاكرة أثناء مرحلة التدريب، باستخدام تنفيذ نموذج Llama-2 على huggingeface.

أود أن أذكر أن هذا الحساب يتم تطبيقه على جميع سلاسل Llama ، وسلسلة Mistral ، وهو قريب من النماذج الأخرى حيث أن جميع النماذج تقريبًا لها نفس البنية، وسوف نقدم تقديرًا لمتطلبات ذاكرة MOE باستخدام تنفيذ نموذج Mixtral في مقال آخر.

متطلبات الذاكرة الإجمالية لأي شبكة عصبية عميقة أثناء وقت التدريب هي كما يلي،

1- ذاكرة لمعلمة النموذج

2- الذاكرة للتدرج "نفس معامل الذاكرة"

3- الذاكرة للمحسن "معامل ذاكرة مضاعف في حالة آدم أو نفس المعاملات في حالة SGD"

4- ذاكرة التنشيط "الجزء الأكبر من استهلاك الذاكرة"

لتسهيل الحساب، سنقوم بحساب هذه العوامل لكل وحدة "MLP والاهتمام" على طبقة فك التشفير

ملف نمذجة اللاما على huggingface هو modeling_llama . سنبدأ بحساب اللاما وفي مقالة مستقبلية حساب MOE .

الصورة السابقة هي هيكل فك التشفير فقط من ورقة GPT ، كما يمكننا أن نرى أن هذه طبقة واحدة من النموذج، و GPT لديه 12 طبقة كما يوضح الشكل.

لذا، يمكننا إجراء الحساب لطبقة واحدة ثم ضرب متطلبات الذاكرة لهذه الطبقة في عدد الطبقات.

سنقوم بفصل طبقة فك التشفير إلى كتلتين رئيسيتين طبقة الاهتمام وطبقة MLP، دعونا نبدأ بطبقة MLP.

قبل الشروع في هذه المهمة الشاقة، دعونا نحدد المعلمة التي نحتاجها لإجراء جميع العمليات الحسابية

حجم الدفعة، طول التسلسل، الحجم المخفي، الطبقات، الحجم المتوسط، الحركة التقدمية، عدد الرؤوس.

آمل أن يكون هذا المعامل معروفًا للقارئ. بالنسبة لـ seq_length، قد يعرفه بعض المستخدمين على أنه طول السياق، وهو عدد الرموز في كل تسلسل، وكما نعلم، هذا هو عنق الزجاجة داخل آلية الانتباه. hidden_size هو المعامل الذي يتحكم في الاتصال المتبقي داخل طبقة فك التشفير. Precession هو متطلب البايت لكل معامل 'float 32(4 byte)، bfloat16 & float16 (2 byte)'. num_heads هو عدد رؤوس الانتباه. يمكنك البحث أكثر عن هذا المصطلح إذا كنت لا تعرف ما يعنيه بالضبط.

طبقة MLP

إذا فتحت هذا الرابط LlamaMLP فسوف تنتقل إلى تنفيذ طبقة MLP من نموذج Llama وستجد أن هذه الوحدة تحتوي على ثلاث طبقات خطية على النحو التالي،

gate_proj (الحجم المخفي، الحجم المتوسط) 
up_proj (الحجم المخفي، الحجم المتوسط) 
down_proj (الحجم المتوسط، الحجم المخفي)

نموذج الذاكرة = 3* (الحجم المخفي * الحجم المتوسط) … (1)

نفترض أن المُحسِّن هو AdamW لجميع العمليات الحسابية، وبالتالي فإن متطلبات الذاكرة للمُحسِّن هي ضعف متطلبات النموذج حيث يحتاج مُحسِّن Adam إلى مقياسين لحساب الزخم. التدرج لهذه الطبقة "كل DNN" هو نفس التدرج في memory_model.

التدرج = نموذج الذاكرة …(2)

المُحسِّن = 2* نموذج الذاكرة … (3)

التنشيط هو ما نحسبه أثناء التمريرة الأمامية من أجل حساب التدرج داخل التمريرة الخلفية ويتم حفظ كل هذا التنشيط في كل تمريرة أمامية، لذلك لدينا مقدار ثابت من الذاكرة في كل تمريرة أمامية. دعنا نحسب ذاكرة التنشيط لـ MLP. إذا أردنا حساب التنشيط لرمز واحد ، فداخل الدالة الأمامية سيكون لدينا ما يلي،

~ hidden_size + medium_size + medium_size = hidden + 2 * medium

تلك الخاصة برمز واحد، إذا أردنا أن نفعل نفس الشيء لتسلسل الرموز، نضرب في طول التسلسل

التنشيط ~= seq_length* hidden_size + 2 * seq_length * medium_size

الشيء الأخير عادةً عندما نقوم بتدريب النموذج نستخدم دفعة صغيرة حتى نضرب المعادلة السابقة بحجم الدفعة.

التنشيط = حجم الدفعة * طول التسلسل * (الحجم المخفي + 2 * الحجم المتوسط) ... (4)

الآن دعونا نلخص كل شيء،

ذاكرة MLP = ((1) + (2) + (3) + (4)) * التقدم

MLP_memory =( 6 * (الحجم المخفي * الحجم المتوسط) + حجم الدفعة * طول التسلسل * (الحجم المخفي + 2 * الحجم المتوسط)) * التقدم

مثال لنموذج Llama 7B،

hidden_size = 4096، inermediate_size = 11008، precession = 2 بايت، seq_length = 2048، batch_size = 8

MLP_memoy ~= 1.3 GB ، باستخدام هذا، نحسب متطلبات الذاكرة لطبقة MLP واحدة.

دعنا ننتقل إلى طبقة الاهتمام.

طبقة الاهتمام

كما في السابق، يمكنك الانتقال إلى تنفيذ طبقة الانتباه في هذا مثل LlamaAttention . كما في السابق، سنقوم بتقسيم الذاكرة الإجمالية إلى نموذج مكون من 4 أقسام، التدرجات، المحسن، التنشيط.

q_proj: (الحجم المخفي، عدد الرؤوس * بُعد الرأس) 
k_proj: (الحجم المخفي، عدد رؤوس القيم الرئيسية * بُعد الرأس) 
v_proj: (الحجم المخفي، عدد رؤوس القيم الرئيسية * بُعد الرأس) 
o_proj: (الحجم المخفي، الحجم المخفي)

هذه هي معلمات الانتباه، بالنسبة لـ num_key_value_heads لن نستخدمها في الحساب من أجل البساطة ولكنها طريقة لتقليل الحساب والذاكرة لمزيد من التفاصيل يمكنك القراءة عن Group Attention. كما فعلنا من قبل، فإن ذاكرة النموذج هي

نموذج الذاكرة = 3 * الحجم المخفي * عدد الرؤوس + الحجم المخفي²…(5)

التدرجات = ذاكرة النموذج …(6)

المُحسِّن = 2* ذاكرة النموذج …(7)

كما ذكرنا أعلاه فإن التنشيط هو المستهلك الأكبر للذاكرة، وخاصة في طبقة الانتباه، إذا ألقيت نظرة على التمريرة الأمامية في وحدة LlamaAttention فسترى هذه المصفوفات الخمس،

حالات الاستعلام :  ( حجم الدفعة ، عدد الرؤوس ، طول  التسلسل ، بُعد الرأس )
 حالات المفتاح :  ( حجم الدفعة ، عدد رؤوس قيمة المفتاح ،  طول التسلسل ، بُعد الرأس )
 حالات القيمة :  ( حجم الدفعة ، عدد رؤوس قيمة المفتاح ،  طول التسلسل ، بُعد الرأس )
 أوزان الانتباه :  ( حجم الدفعة ، عدد الرؤوس ،  طول التسلسل ،  طول التسلسل )
 إخراج الانتباه :  ( حجم الدفعة ،  طول التسلسل ، الحجم المخفي )

head_dim = hidden_size // عدد الرؤوس

إجمالي متطلبات الذاكرة للتنشيط هي،

التنشيط = حجم الدفعة * ( 3 * عدد الرؤوس * طول التسلسل * أبعاد الرأس + عدد الرؤوس * طول التسلسل² + طول التسلسل * الحجم المخفي)

ذاكرة الاهتمام = ( 6 * (الحجم المخفي * عدد الرؤوس + الحجم المخفي²) + حجم الدفعة * ( 3 * عدد الرؤوس * طول التسلسل * بُعد الرأس + عدد الرؤوس * طول التسلسل² + طول التسلسل * الحجم المخفي))* التقدم

نموذج فك التشفير فقط

جميع الحسابات السابقة خاصة بطبقة واحدة، لذا لحساب متطلبات الذاكرة الكاملة، تحتاج إلى الضرب في عدد الطبقات

إجمالي ذاكرة النموذج = عدد الطبقات *(ذاكرة MLP + ذاكرة الاهتمام)

هذه هي الذاكرة التي تحتاجها لتدريب وضع فك التشفير فقط، لقد قمت بإنشاء مساحة llm-memory-requiremen على huggingface يمكنك إلقاء نظرة عليها واللعب بالمعلمات لتقدير ما تحتاجه لتدريب نموذج باستخدام تكوينات مختلفة.

خاتمة

في الختام، يوضح هذا المقال كيفية تقدير متطلبات الذاكرة لنموذج فك التشفير فقط، ونستخدم تنفيذ Llama لحساب ذلك، ويمكننا ملاحظة التكلفة الإضافية للذاكرة من حساب الانتباه والتي تكون تربيعية مع seq_length.

خليل حنارة
مهندس ذكاء اصطناعي في مسراج